#!/usr/bin/python
# Copyright 1999-2014 Gentoo Foundation
# Distributed under the terms of the GNU General Public License v2
# $Header: $

"""Convert a Gentoo system using SYMLINK_LIB=yes to SYMLINK_LIB=no

This script assumes that everything will Just Work, so there is little error
checking whatsoever included.  It is also a work in progress, so there may be
bugs.  Please report any bugs found to http://bugs.gentoo.org/.
"""

from __future__ import print_function

import argparse
import datetime
import errno
import os
import subprocess
import sys


WARNING_INPUT = 'I understand this may break my system and I have a backup'
WARNING = """
Please enter the following sentence (with punctation and capitalization) and
press Enter, or press Ctrl-C to quit:"""


def lutimes(path, times):
	d = datetime.datetime.fromtimestamp(int(times)).isoformat()
	subprocess.check_call(['touch', '-h', '-d', d, path])


class Entry(object):
	"""Object to hold a single line of a CONTENTS file"""

	def __init__(self, line):
		line = line.rstrip('\n')
		self.type, line = line.split(' ', 1)
		if self.type == 'dir':
			self.path = line
		elif self.type == 'obj':
			self.path, self.hash, self.time = line.rsplit(' ', 2)
		elif self.type == 'sym':
			line, self.time = line.rsplit(' ', 1)
			self.path, self.target = line.split(' -> ')
		else:
			raise ValueError('cannot handle %s %s' % (self.type, line))

	def __str__(self):
		eles = [self.type]
		if self.type == 'dir':
			eles.append(self.path)
		elif self.type == 'obj':
			eles.append(self.path)
			eles.append(self.hash)
			eles.append(self.time)
		elif self.type == 'sym':
			eles.append(self.path)
			eles.append('->')
			eles.append(self.target)
			eles.append(self.time)
		return ' '.join(eles)


def atomic_write(path, content):
	"""Write out a new file safely & atomically"""
	new_path = '%s.new' % path
	with open(new_path, 'w') as f:
		f.write(content)
	# We don't worry about privacy here as all the files we're updating
	# are world readable and lack secrets.
	st = os.lstat(path)
	lutimes(path, st.st_mtime)
	os.chown(new_path, st.st_uid, st.st_gid)
	os.chmod(new_path, st.st_mode)
	os.rename(new_path, path)


def convert_root(root, dry_run=True, verbose=False):
	"""Convert all the symlink paths in |root|"""

	# Set SYMLINK_LIB=no if need be.
	make_conf = root + 'etc/portage/make.conf'
	if os.path.exists(make_conf):
		with open(make_conf) as f:
			content = f.read()
		if 'SYMLINK_LIB=no' not in content:
			print('Setting SYMLINK_LIB=no in %s ...' % make_conf)
			if not dry_run:
				atomic_write(make_conf, content + '\n'.join([
					'',
					'# START: AUTO-UPGRADE SECTION',
					'# Remove these lines after upgrading your profile to 14.0+.',
					'SYMLINK_LIB=no',
					'LIBDIR_x86=lib',
					'LIBDIR_ppc=lib',
					'# END: AUTO-UPGRADE SECTION',
					'',
				]))
	else:
		print('Please set SYMLINK_LIB=no in your package manager config files')

	# First make sure the various lib paths are not symlinks.
	libs = ['lib', 'usr/lib', 'usr/local/lib']
	for p in libs:
		rp = os.path.join(root, p)
		t = None
		if os.path.islink(rp):
			print('Removing %s symlink ...' % rp)
			t = os.lstat(rp).st_mtime
			if not dry_run:
				os.unlink(rp)

			rp32 = rp + '32'
			if os.path.isdir(rp32):
				print('Renaming %s to %s ...' % (rp32, rp))
				if not dry_run:
					os.rename(rp32, rp)

		if not os.path.exists(rp):
			print('Creating %s dir ...' % rp)
			if not dry_run:
				os.makedirs(rp)

		if t:
			if not dry_run:
				os.utime(rp, (t, t))

	show = {
		'cat': False,
		'pkg': False,
	}
	def showit(cat, pkg):
		if not show['cat']:
			show['cat'] = True
			print('Processing category %s ...' % cat)
		if pkg is None:
			return

		if not show['pkg']:
			show['pkg'] = True
			print('  %s ...' % pkg)

	# Now walk the vdb looking for files that installed into /lib32 and /lib.
	vdb = os.path.join(root, 'var', 'db', 'pkg')
	for cat in os.listdir(vdb):
		vdb_cat = os.path.join(vdb, cat)
		if not os.path.isdir(vdb_cat):
			continue

		show['cat'] = False
		if verbose:
			showit(cat, None)

		for pkg in os.listdir(vdb_cat):
			vdb_pkg = os.path.join(vdb_cat, pkg)

			show['pkg'] = False
			if verbose:
				showit(cat, pkg)

			contents = os.path.join(vdb_pkg, 'CONTENTS')
			if not os.path.exists(contents):
				if verbose:
					print('SKIP')
				continue

			# Process the package's contents and rename things as needed.
			modified = False
			new_contents = []
			with open(contents) as f:
				for line in f:
					e = Entry(line)

					if e.type == 'obj' or e.type == 'sym':
						# Migrate files from /usr/lib64/ that really belong in /usr/lib/.
						for p in libs:
							p = '/%s' % p
							if not e.path.startswith(p + '/'):
								continue

							src = '%s64/%s' % (p, e.path[len(p) + 1:])
							# Make sure the source still exists.  Maybe it was migrated
							# already or the user somehow deleted it.
							rs = os.path.normpath(root + src)
							if not os.path.lexists(rs):
								continue

							rd = os.path.normpath(root + e.path)
							# Make sure the destination doesn't exist.  This could happen
							# when a /lib32/foo moved to /lib/foo.
							if not dry_run:
								if os.path.exists(rd):
									continue

							try:
								if not dry_run:
									os.makedirs(os.path.dirname(rd))
							except OSError as ex:
								if ex.errno != errno.EEXIST:
									raise

							showit(cat, pkg)
							if os.path.islink(rd) and not os.path.islink(rs):
								print('    SKIP %s' % e.path)
								continue

							print('    MOVE %s -> %s' % (src, e.path))

							if not dry_run:
								os.rename(rs, rd)
							# Clean up empty dirs.
							try:
								if not dry_run:
									while True:
										rs = os.path.dirname(rs)
										os.rmdir(rs)
							except OSError as ex:
								if ex.errno != errno.ENOTEMPTY:
									raise

					if e.type == 'dir' or e.type == 'obj' or e.type == 'sym':
						# Update the location of files in /lib32/ to /lib/.
						for p in libs:
							p = '/%s' % p
							p32 = '%s32' % p
							if e.path == p32:
								new_path = p
							elif e.path.startswith(p32 + '/'):
								new_path = '%s/%s' % (p, e.path[len(p32) + 1:])
							else:
								continue
							showit(cat, pkg)
							print('    CONT %s -> %s' % (e.path, new_path))
							e.path = new_path

					if e.type == 'sym':
						# Handle symlinks that point to files in /lib32/.
						for p in libs:
							p = '/%s' % p
							p32 = '%s32' % p
							if p32 in e.target:
								new_path = e.target.replace(p32, p)
								showit(cat, pkg)
								print('    LINK %s -> %s' % (e.target, new_path))
								e.target = new_path

								rl = os.path.normpath(root + e.path)
								if os.path.islink(rl):
									if not dry_run:
										os.unlink(rl)
										os.symlink(e.target, rl)
										lutimes(rl, e.time)

						# Handle symlinks for the dirs themselves.  e.g. glibc
						# will create /lib -> lib64
						for p in libs:
							if e.path == '/%s' % p:
								e = None
								break

					if e:
						e = str(e)
						if e != line.rstrip('\n'):
							modified = True
						new_contents.append(e)
					else:
						modified = True

			# Now write out the new CONTENTS.
			if not dry_run and modified:
				content = '\n'.join(new_contents)
				if content:
					content += '\n'
				atomic_write(contents, content)


def main(argv):
	# Process user args.
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument('--root', type=str, default=os.environ.get('ROOT', '/'),
	                    help='ROOT to operate on')
	parser.add_argument('-n', '--dry-run', default=True, action='store_true',
	                    help='Do not make any changes')
	parser.add_argument('--wet-run', dest='dry_run', action='store_false',
	                    help='Make changes to the filesystem')
	parser.add_argument('-v', '--verbose', default=False, action='store_true',
	                    help='Show all packages checked')
	opts = parser.parse_args(argv)

	if not os.path.isdir(opts.root):
		parser.error('root "%s" does not exist' % opts.root)
	opts.root = os.path.normpath(opts.root).rstrip('/') + '/'

	# Verify the user wants to run us.
	if not opts.dry_run:
		try:
			input = raw_input('%s\nWill operate on ROOT=%s\n%s\n\n%s\n\n-> ' %
			                  (__doc__, opts.root, WARNING, WARNING_INPUT))
			input = WARNING_INPUT
		except (EOFError, KeyboardInterrupt):
			input = None
		if input != WARNING_INPUT:
			print('\nAborting...')
			return os.EX_USAGE

	# Let's gogogogogo.
	print('Checking system for old lib32 dirs')
	convert_root(opts.root, opts.dry_run, opts.verbose)

	# Run some checkers after the fact.
	try:
		print('\nRunning qcheck on your system; '
		      'you might want to re-emerge any broken packages')
		if opts.dry_run:
			print(' ... skipping checks ...')
		else:
			subprocess.check_call(['qcheck', '-aB', '--root', opts.root])
		print(' ... No broken packages! woot!')
	except subprocess.CalledProcessError:
		pass

	print('\nAll finished!')


if __name__ == '__main__':
	sys.exit(main(sys.argv[1:]))