#!/usr/bin/perl -w

# autochroot
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; version 2 of the License.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

use strict;
use Cwd;
use File::Copy;
use File::Temp qw/ tempfile tempdir /;

#
# declarations
#
sub message($);
sub warning($);
sub error($);
sub debug($);
sub trim($);
sub my_color($);
sub my_mount($$);
sub my_mount_opt($$$);
sub my_umount();
sub cleanup($);
sub sighandler();
sub parseInput($@);
sub choose_distro();
sub getversion($);
sub arch_compatible($);
sub do_chroot($);
sub chroot_devnode($);
sub architectureMatch($);
sub usage();

#
# install sighandlers (to be able to clean up)
#
$SIG{INT} = \&sighandler;
$SIG{TERM} = \&sighandler;

#
# globals
#
our ($MOUNTPOINT, $TARGET, $DEBUG, $COLORS, @MOUNTS, $TMPDIR, %CONFIGFILES);
$MOUNTPOINT = "/mnt";

$DEBUG = 1;

# import Term::ANSIColor if existing (or disable colored output for my_color())
eval("use Term::ANSIColor");
if ($@) {
	$COLORS = 0;
	warning("Couldn't load terminal color module. Continuing without colors");
}
else {
	$COLORS = 1;
}


#
# trim function
#
# Removes whitespaces at the beginning and the end of the input string and returns the result
# Parameters:
#   string: string to trim
# Returns:
#   trimmed result string
#
sub trim($) {
	my $string = shift;
	$string =~ s/^\s+//;
	$string =~ s/\s+$//;
	return $string;
}

#
# write a red error message to STDERR
#
# Parameters:
#   msg: Error message
sub error($) {
	my $msg = shift;
	print STDERR my_color('bold red')."Error: $msg\n".my_color('reset');
}

#
# write a green warning message to STDERR
#
# Parameters:
#   msg: Warning message
sub warning($) {
	my $msg = shift;
	print STDERR my_color('green')."Warning: $msg\n".my_color('reset');
}

#
# print a yellow message to STDOUT
#
# Parameters:
#   msg: Message
sub message($) {
	print my_color('yellow').(shift).my_color('reset');
}

#
# print a cyan debug message to STDERR if the global variable $DEBUG is true
#
# Parameters:
#   text: Debug message
sub debug($) {
	return 0 if not $DEBUG;

	my $text = shift;

	if ($text) {
		print STDERR my_color('cyan')."Debug: $text\n".my_color('reset');
	}
}

#
# thin wrapper for Term::ANSIColor's color() that returns nothing if the color mode is disabled
# Parameters:
#   col: new color
# Returns:
#   color escape sequence or '' if colored output is disabled
sub my_color($) {
	my $rc = '';

	if ($COLORS) {
		my $col = shift;
		$rc = eval("color('$col')");
	}

	return $rc;
}

#
# Thin wrapper for my_mount_opt($$$) that defaults the options (third parameter) to ''
# Parameters:
#   dev: Device node
#   mp:  mount point
# Returns:
#   1 if successful, 0 otherwise
sub my_mount($$) {
	return my_mount_opt(shift, shift, '');
}

#
# Mounts a filesystem and adds the mount point to the mount list (to be unmounted by cleanup())
# Parameters:
#   dev:     Device node
#   mp:      mount point
#   options: mount options (including '-o ') 
# Returns:
#   1 if successful, 0 otherwise
sub my_mount_opt($$$) {
	my $dev = shift;
	my $mp = shift;
	my $options = shift;
	my $rc = 0;

	debug "Mounting '$dev' at '$mp'";
	# check if anything's already mounted at $mp
	if (system("mountpoint -q $mp")) {
		if (!system("mount $options $dev $mp")) {
			$rc = 1;
			$MOUNTS[@MOUNTS] = $mp;
		}
	}
	else {warning "'$mp' already is a mountpoint!\n";}

	return $rc;
}

#
# unmount all previously mounted file systems in reverse order and clear @MOUNTS
# when done
#
sub my_umount() {
	my $mp;

	@MOUNTS = reverse(@MOUNTS);
	foreach $mp (@MOUNTS) {
		debug "Unmounting '$mp'...";
		# repeatedly unmount the partition (just to be sure)
		while (!system("umount $mp 2>/dev/null")) {}
	}
	@MOUNTS = ();
	return 1;
}

#
# cleanup all the allocated resources and exit with the given exit code
#
# Parameters:
#   rc: exit code
# Returns:
#   This function does not return!
#
sub cleanup($) {
	my $rc = shift;

	# restore config files
	my @confs = keys %CONFIGFILES;
	foreach my $c (@confs) {
		my $tmp = $CONFIGFILES{$c};
		debug "restoring '$c' using '$tmp'...";
		
		# copy backup file over modified config
		if (copy($tmp, $c)) {
			unlink($tmp) or warning ("Couldn't remove temporary file '$tmp'");
		}
		else {
			warning "Couldn't restore config file '$c'!";
		}
	}
	%CONFIGFILES = ();

	# unmount
	my_umount();

	# remove temporary mount directory
	if ($TMPDIR) {
		rmdir $TMPDIR;
	}

	exit $rc;
}

#
# Signal handler
#
# Catches SIGINT and SIGTERM and calls cleanup()
#
sub sighandler() {
	error("Signal caugth! exiting!");
	cleanup(1);
}

#
# parse the user input after the distribution list was printed
#
# Check if the user entered a device node, a mountpoint or a list entry's index
# (or wants to quit autochroot)
#
# Parameters:
#   devices: reference to the distributions list's device nodes
# Returns:
#   1 on success, 0 on failure
#
sub parseInput($@) {
	my $in = shift;
	my $ref_devices = shift;
	my @devices = @{$ref_devices};
	my $rc = 0;

	# empty input => select first distro
	$in = 1 if ($in eq '');

	# quit
	if ($in eq '0' || $in eq 'q') {
		error("Script aborted\n");
		return 1;
	}
	# list entry
	elsif ($in =~ m/[0-9]+/) {
		if ($in le @devices) {
			$rc = chroot_devnode($devices[$in-1]);
		}
		else {
			error("Entry not found!");
		}
	}
	# device node
	elsif ($in =~ m/^\/dev/) {
		$rc = chroot_devnode($in)
	}
	# absolute path
	elsif ($in =~ m/^\//) {
		$rc = do_chroot($in);
	}
	# relative path
	elsif (-d $in ) {
		$rc = do_chroot($in);
	}
	else {
		error('Unexpected input');
	}

	return $rc;
}

#
# find device nodes that could contain linux distributions, mount them if necessary
# and try to find version information.
# Then print all found distributions with version and hostname information
#
sub choose_distro() {
	#get all ext-filesystems
	my @ext_devices=split(/\n/, `/sbin/fdisk -l|grep "Linux\$" |cut '-d ' -f1`);
	my ($dev, $mp, $um, $vers, $im, $compatible, $i);
	my (@devices, @mountpoints, @distros, @comp);
	my $rc = 0;
	
	# check each found ext-partition
	foreach $dev (@ext_devices) {
		# get the mountpoint if the device is already mounted
		$mp=trim(`grep ^$dev /etc/mtab|head -n1|cut '-d ' -f2`);

		# umount flag (specifies if the partition has been mounted by autochroot)
		$um = 1;
		if (!$mp) { # partition isn't mounted yet => mount it
			if (my_mount($dev, $MOUNTPOINT)) {
				$um = 1;
				$mp = $MOUNTPOINT;
			}
		}

		$vers = '';
		$compatible = '';
		if ($mp ne "/") { # skip the current root file system (don't show it in the distro list)
			$vers = getversion($mp);
			if ($vers) {
				$compatible = arch_compatible($mp); #check compatibility
			}
		}

		if ($vers) { # found version info => add to distro list
			# TODO: use a nested hashmap
			$devices[@devices] = $dev;
			$mountpoints[@mountpoints] = $um ? '' : $mp;
			$distros[@distros] = $vers;
			$comp[@comp] = $compatible
		}

		# unmount partitions mounted by autochroot
		my_umount();
	}

	# repeatedly read user input until it's valid (parseInput() returns !=0)
	while (!$rc) {
		if (@devices) {
			print "Found distributions:\n";
			my $len = int(log(@devices)/log(10))+1;
			
			# print distribution list
			for ($i=0; $i < @devices; $i++) {
				printf('  (%'.$len.'i) ', $i+1);
				print my_color('green').$devices[$i].my_color('reset');
				if ($mountpoints[$i]) {
					print ' at ';
					print my_color('green').$mountpoints[$i].my_color('reset');
				}
				print " $distros[$i]";
				if ($comp[$i] == 0) {
					print my_color('bold red').' INCOMPATIBLE'.my_color('reset');
				}
				print("\n");
			}
		}
		else {
			error('No Distributions found!');
		}

		print "\nYou've got the following options:\n";
		if (@devices) { #show this option only if distributions were found
			print "- Enter the number of the desired distribution\n"
				."- Simply press return (to choose the first list item)\n";
		}
		print "- Enter a device node directly (e.g. /dev/sda1)\n"
			."- Enter a mountpoint (e.g. /mnt/foo or /media/foo)\n"
			."- Enter '0' or 'q' to quit\n"
			.":: ";
		my $in = <STDIN>;
		if (not $in) { # CTRL+D pressed => not even a \n in STDIN
			error("Script aborted\n");
			return 1;
		}
		$in = trim($in);

		$rc = parseInput($in, \@devices);
	}
	return $rc;
}

#
# search for version information within a given mountpoint
# Parameters:
#   mp: mount point
# Returns:
#   Version information
#
sub getversion($) {
	my $mp=shift;
	my ($dname, $vfile);
	my $rc = 0;

	#versionfiles is based on infobash's list (contact me if something's wrong/missing)
	my %versionfiles = (
		'/etc/debian_version'		=> 'Debian',
		'/etc/gentoo-release'		=> 'Gentoo',
		'/etc/mandrake-release'		=> 'Mandrake',
		'/etc/slackware-version'	=> 'Slackware',
		'/etc/redhat-release'		=> 'RedHat',
		'/etc/SuSE-release'		=> 'SuSE',
		'/etc/sidux-version'		=> 'Sidux',
		'/etc/ubuntu_version'		=> 'Ubuntu',
		'/etc/kanotix-version'		=> 'Kanotix',
		'/etc/knoppix-version'		=> 'Knoppix',
		'/etc/sabayon-release'		=> 'Sabayon',
		'/etc/turbolinux-release'	=> 'TurboLinux',
		'/etc/zenwalk-version'		=> 'Zenwalk'
		);

	foreach $vfile (keys %versionfiles) {
		$dname = $versionfiles{$vfile};
		if (-e "$mp/$vfile") {
			if (open(VHANDLE, "$mp/$vfile")) {
				$rc="$dname: ".trim((<VHANDLE>)[0]);
				
				#if there is a /etc/hostname, read it and append it to the output
				if (open(VHANDLE2, "$mp/etc/hostname")) {
					$rc .= ' (Hostname: '.trim((<VHANDLE2>)[0]).')';
					close(VHANDLE2);
				}
				close(VHANDLE);
			}
		}
	}
	return $rc;
}

#
# check compatibility of the distribution architectures
# Parameters:
#   mp: mount point
# Returns:
#   1 if compatible or unknown, 0: incompatible (e.g. 64bit guest on a 32bit kernel)
#
sub arch_compatible($) {
	my $rc = 0;
	my $mp = shift;
	my $arch = trim(qx(uname -m));
	my $myarch;

	if (! -f '/usr/bin/readelf') {
		debug "readelf not found, can't check compatibilities";
		return 1
	}

	# get my own architecture
	if ($arch eq 'x86_64' || $arch eq 'ia64') {
		$myarch = 'x64';
	}
	elsif ($arch =~ m/i.86/) {
		$myarch = 'x86';
	}
	else {
		debug "unknown architecture: $arch";
		return 1;
	}

	#get $mp architecture
	$arch = trim(qx(/usr/bin/readelf -h $mp/bin/bash|grep Class|cut -d: -f2));

	if ($myarch eq 'x64' && ($arch eq 'ELF32' || $arch eq 'ELF64')) { 
		$rc = 1;
	}
	elsif ($myarch eq 'x86' && $arch eq 'ELF32') {
		$rc = 1;
	}

	return $rc;
}

#
# chroot to a devnode
#
# Check if the given devnode is mounted (mount it if it isn't) and call do_chroot()
# Parameters:
#   dn: Device node
# Returns:
#   1
#
sub chroot_devnode($) {
	my $dn = shift;
	my $mp;

	# check if it's mounted
	$mp=trim(`grep ^$dn /etc/mtab|head -n1|cut '-d ' -f2`) if ($dn);

	if (!$mp) { # mount it (and mark it for unmounting)
		if (my_mount($dn, $MOUNTPOINT)) {
			$mp = $MOUNTPOINT;
		}
		
	}

	if ($mp) {
		do_chroot($mp);
	}
	else {
		error("couldn't get a mountpoint for the device node!");
	}

	my_umount();

	return 1;
}

#
# mount pseudo-filesystems and perform the chroot
# Parameters:
#   mp: Mount point
# Returns:
#   1 on success, 0 on failure
#
sub do_chroot($) {
	my $mp = shift;
	my $rc = 0;

	if ($mp ne '/') {
		my $err = 0;
		print ("mounting necessary filesystems in '$mp'...\n");
		if (! -d "$mp/dev") {
			error("No /dev directory found on the partition!");
			$err = 1;
		}
		if (! -d "$mp/proc") {
			error("No /proc directory found on the partition!");
			$err = 1;
		}
		if (! -d "$mp/sys") {
			error("No /sys directory found on the partition!");
			$err = 1;
		}

		if (!$err) {
			if (!my_mount_opt('/dev/', "$mp/dev", '--bind')) {
				error("Couldn't mount /dev");
				$err = 1;
			}
			if (!my_mount_opt('none', "$mp/proc", '-t proc')) {
				error("Couldn't mount /proc");
				$err = 1;
			}
			if (!my_mount_opt('none', "$mp/sys", '-t sysfs')) {
				error("Couldn't mount /sys");
				$err = 1;
			}
			
			copyConfigfile('/etc/resolv.conf', "$mp/etc/resolv.conf");

			if (!$err) {
				print ("chrooting to '$mp'\n");
				
				my $pid = fork();
				if (!$pid) {
					#child
					if (chdir($mp)) {
						my $cwd = getcwd();
						# copy mtab (mount table)
						if (system("grep '$cwd' /etc/mtab|grep -v bind|sed 's#$cwd#/#g'|sed 's#//#/#g' > 'etc/mtab'")) {
							warning("Couldn't copy mtab! You won't be able to resolve mount points within the jail!\n");
						}

						if (chroot($cwd)) {
							message("  To exit the chroot jail press CTRL+D or enter exit\n");
							exec('/bin/bash');
						}
					}
					error("couldn't chroot!\n");
					cleanup(-1); # 
				}
				elsif ($pid>0) {
					# parent
					wait();
					$rc = 1;
				}
				else {
					error("Error: Couldn't create a child process!\n");
				}

				print ("cleaning up...\n");
				my_umount();
			}
		}
	}
	else {
		error("Can't chroot to '/'!");
	}
	return $rc;
}


#
# backup the original config file and copy another one over it
# Parameters:
#   src: source file
#   tgt: target file (will be backed up and then overwritten)
# Returns:
#   1: successful, 0: failure
#
# TODO: remove the copy part and rename me to backupConfig($)
sub copyConfigfile($$) {
	my $src = shift;
	my $tgt = shift;
	my $rc = 0;

	# create backup file
	(my $handle, my $tmp)= tempfile('autochroot.conf.XXXXXX', DIR=>'/tmp');

	debug ("Copying config file '$src' to '$tgt' (backup: '$tmp')");
	
	if ($tmp) {
		# revoke permissions to the backup file
		chmod 0600, $tmp;
		
		if (-f $tgt) {
			#add to %CONFIGFILES
			$CONFIGFILES{$tgt} = $tmp;
		}
		else {
			# original file doesn't exist => remove temporary file
			unlink $tmp;
		}

		# backup (if original config file exists)
		if (!-f $tgt || copy ($tgt, $tmp)) {
			# overwrite original file
			if (copy ($src, $tgt)) {
				$rc = 1;
			}
			else {
				warning("Couldn't copy config file '$src'!\n");
			}
		}
		else {
			warning("Couldn't create backup of '$src'!\n");
		}
	}
	
	return $rc;
}

#
# print usage information and exit
#
sub usage() {
	print my_color('bold')."USAGE:\n".my_color('reset');
	print "\t$0 [options] [target]\n"
		. "\t   If you don't specify a ".my_color('bold')."target".my_color('reset').", the script will run interactively.\n"
		. my_color('bold')."OPTIONS:\n".my_color('reset')
		. "\t--mountpoint, -m (default: '/mnt/autochroot.XXXXXX' with 'XXXXXX' to be\n"
		. "\t  a random string, read man mkstemp for details)\n"
		. "\t  set the path the specified device node should be mounted at\n"
		. "\t  (this option is only used when you specify a device node as [target]\n"
		. "\t  or interactively)\n"
		. "\n"
		. my_color('bold').'NOTES:'.my_color('reset')."\n"
		. "\t$0 has to be run as root user\n";
	cleanup 1;
}

#
# check for root
#
if ($< != 0) {
	error("This script has to be run as root!");
	cleanup 1;
}


$TARGET = '';

while (@ARGV) {
	my $param = shift(@ARGV);

	# change mount point (used only in combination with dev nodes or in interactive mode)
	if ($param eq '--mountpoint' || $param eq '-m') {
		$MOUNTPOINT=shift(@ARGV);
		debug "new mountpoint: $MOUNTPOINT";
	}
	# print help
	elsif ($param eq '--help' || $param eq '-h') {
		usage();
	}
	# new target
	elsif ($param =~ m/^[^-]/) {
		$TARGET = $param;
		debug "new target: $TARGET";
	}
}

#
# create temporary directory to mount filesystems
#

$TMPDIR = tempdir('autochroot.XXXXXX', DIR => '/mnt');
if ($TMPDIR) {
	debug "temporary mount directory: $TMPDIR";
	$MOUNTPOINT = $TMPDIR;
}
else {
	error("Couldn't create temporary directory to mount filesystems, aborting\n");
	cleanup 1;
}

my $rc = 0;

if ($TARGET ne '') {
	my @foo = ();
	$rc = parseInput($TARGET, \@foo);
}
else {
	$rc = choose_distro();
}

cleanup $rc;