#!/usr/bin/python2.7
# -*- coding: iso-8859-1 -*-

import argparse, os, socket, subprocess, sys

# read arguments
parser = argparse.ArgumentParser(description='Enter chroot build environment for a host')
parser.add_argument('-H', '--host')
parser.add_argument('-k', '--kernel', action='store_true')
parser.add_argument('-c' , '--command')
parser.add_argument('-e', '--exit', action='store_true')
parser.add_argument('-w', '--world', action='store_true')
parser.add_argument('-v', '--verbose', action='store_true')
args = parser.parse_args()

# must be run as root
if os.getuid() != 0:
	print '\n%smust be run as root!\n' % os.path.basename(sys.argv[0])
	parser.print_help()
	sys.exit(1)

# get host from name or command line
if os.path.basename(sys.argv[0]) != 'ch':
	args.host = os.path.basename(sys.argv[0])[2:]

if not args.host:
	parser.print_help()
	sys.exit(1)

if args.verbose:
	print 'Creating autorun file'

# create .autorun file if asked
if args.world:
	args.command = 'world-update'
	args.exit = True

if args.command:
	arun = open('/mnt/%s/.autorun' % args.host, 'w')
	arun.write('%s\n' % args.command)
	if args.exit:
		arun.write('exit\n')
	arun.close()
else:
	try:
		os.remove('/mnt/%s/.autorun' % args.host)
	except OSError:
		pass

if args.verbose:
	print 'Pinging host'

# check if host is available for syncing
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.settimeout(10)
try:
	s.connect((args.host, 22))
	s.close()
	sync_paths = ['etc/portage/', 'var/lib/portage/world', 'var/lib/portage/world_sets', 'usr/local/bin/', 'home/nelz/bin/']
	sync_cmd = ['rsync', '--archive', '--delete-before']

	if args.verbose:
		print 'Syncing host'
		sync_cmd.append('--verbose')

	if args.kernel:
		print 'Syncing kernel sources... be patient'
		sync_paths.extend(['usr/src/', 'lib/modules/', 'boot/'])

	for path in sync_paths:
		if os.path.exists(os.path.join('/mnt/', args.host, path)):
			src_dest = ['%s:/%s' % (args.host, path), os.path.join('/mnt/', args.host, path)]
			if args.verbose:
				print 'syncing', ' to '.join(src_dest)
			subprocess.call(sync_cmd + src_dest)

except socket.error:
	print args.host, 'not available, syncing skipped'

if args.verbose:
	print 'Setting up chroot mounts'

# set up mounts in the chroot
try:
	os.mkdir('/var/tmp/%s' % args.host, 1777)
except OSError:
	pass

subprocess.call(['mount', '-t', 'proc', 'none', '/mnt/%s/proc' % args.host])
subprocess.call(['mount', '--bind', '/var/tmp/%s' % args.host, '/mnt/%s/var/tmp' % args.host])
for dir in ('dev', 'var/portage', 'mnt/portage'):
	subprocess.call(['mount', '--bind', '/%s' % dir, '/mnt/%s/%s' % (args.host, dir)])
subprocess.call(['cp', '-f', '/etc/resolv.conf', '/mnt/%s/etc/' % args.host])

if args.verbose:
	print 'Entering chroot'

# enter chroot
if os.path.exists('/mnt/%s/bin/zsh' % args.host):
	#subprocess.call(['chroot', '/mnt/%s' % args.host, '/bin/zsh'])
	subprocess.call(['su' , '-', '-c', 'chroot /mnt/%s /bin/zsh' % args.host])
else:
	#subprocess.call(['chroot', '/mnt/%s' % args.host, '/bin/bash'])
	subprocess.call(['su' , '-', '-c', 'chroot /mnt/%s /bin/bash' % args.host])

if args.verbose:
	print 'Cleaning up'

# clean up after exiting chroot
try:
	os.remove('/mnt/%s/.autorun' % args.host)
except OSError:
	pass

for dir in ('var/tmp', 'mnt/portage', 'var/portage', 'proc', 'dev'):
	subprocess.call(['umount' , '/mnt/%s/%s' % (args.host, dir)])