#!/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):
  """Change the time stamp on a symlink"""
  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)

          # The gcc specs expect lib32 when using -m32 until it gets rebuilt.
          print('Creating compat symlink %s -> lib ...' % rp32)
          os.symlink('lib', rp32)

    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('--no-prompt', dest='prompt', default=True,
                      action='store_false',
                      help='Assume you know what you are doing')
  parser.add_argument('-v', '--verbose', default=False, action='store_true',
                      help='Show all packages checked')
  parser.add_argument('--no-post-checks', dest='post_checks', default=True,
                      action='store_false',
                      help='Do not run checks after converting everything')
  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 and opts.prompt:
    try:
      resp = raw_input('%s\nWill operate on ROOT=%s\n%s\n\n%s\n\n-> ' %
                       (__doc__, opts.root, WARNING, WARNING_INPUT))
    except (EOFError, KeyboardInterrupt):
      resp = None
    if resp != 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.
  if opts.post_checks:
    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!\nYou should re-emerge glibc, gcc, and binutils.')


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