#!/usr/bin/python
# Copyright 2014 The Chromium OS Authors. All rights reserved.
# 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/.

Things that are known to not work (and cannot be handled):
 - Tools that create files/symlinks in pkg_* funcs or otherwise outside of
   the src_install func.  From the PM's perspective, no one owns these files,
   so it's not really possible to migrate them.
"""

from __future__ import print_function

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


# Python 3 glue.
if sys.hexversion >= 0x3000000:
  # pylint: disable=W0622
  raw_input = input


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

LIBDIRS = ('lib', 'usr/lib', 'usr/local/lib')


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

  Args:
    path: The file to write |content| too; must already exist.
    content: The data to write to the new file.
  """
  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 set_symlink_lib_no(root, dry_run=True,
                       verbose=False): # pylint: disable=W0613
  """Set SYMLINK_LIB=no in the system's make.conf if needed

  If the system already have SYMLINK_LIB=no set, then we don't do anything.

  We also assume that no one has already set this, so the level of checks
  here is fairly low.  For example, if you have:
    #SYMLINK_LIB=no
  We do not detect that.
  """
  # 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')


def move_libdirs(root, dry_run=True, verbose=False): # pylint: disable=W0613
  """Convert lib symlink to a dir, and lib32 dir to a symlink

  Just rename the base paths as needed before doing anything else.
  """
  # First make sure the various lib paths are not symlinks.
  for p in LIBDIRS:
    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 and not dry_run:
      os.utime(rp, (t, t))


def showit(show, cat, pkg):
  """Display status messages as needed"""
  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)


def orphaned_cleanup(path):
  """Deal with possible orphaned files related to this path"""
  # We pre-compile pyc/pyo files in the pkg_* stages, so clean those up.
  if path.endswith('.py'):
    # Various versions of python have generated files like:
    # SlotObject.py  (original file)
    # SlotObject.pyc
    # SlotObject.pyo
    for sfx in ('c', 'o'):
      if os.path.exists(path + sfx):
        print('    CLEAN %s' % (path + sfx))
        os.unlink(path + sfx)

    # While others have used subdirs:
    # __pycache__/SlotObject.cpython-32.pyc
    # __pycache__/SlotObject.cpython-32.pyo
    # Prune the __pycache__ directory entirely as it should only contain
    # cache files.
    cachedir = os.path.join(os.path.dirname(path), '__pycache__')
    if os.path.exists(cachedir):
      print('    CLEAN %s/*' % cachedir)
      shutil.rmtree(cachedir)


def migrate_one_entry(show, cat, pkg, e, root, dry_run=True,
                      verbose=False): # pylint: disable=W0613
  """Process a single entry of the CONTENTS file"""
  # Handle common issues with plain files and symlinks.
  if e.type == 'obj' or e.type == 'sym':
    # Migrate files from /usr/lib64/ that really belong in /usr/lib/.
    # Since /usr/lib was a symlink to /usr/lib64 in the rootfs, while
    # packages installed files into both dirs, the merge process ends
    # up writing them all to /usr/lib64.  We need to walk the contents
    # and file all paths registered in /usr/lib and move them out of
    # /usr/lib64 and back into /usr/lib.
    for p in LIBDIRS:
      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 and 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(show, 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)
        orphaned_cleanup(rs)

      # 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

  # Handle common issues with all entry types.
  if e.type == 'dir' or e.type == 'obj' or e.type == 'sym':
    # Update the location of files in /lib32/ to /lib/.
    for p in LIBDIRS:
      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(show, cat, pkg)
      print('    CONT %s -> %s' % (e.path, new_path))
      e.path = new_path

  # Handle issues specific to symlinks.
  if e.type == 'sym':
    # Handle symlinks that point to files in /lib32/.
    for p in LIBDIRS:
      p = '/%s' % p
      p32 = '%s32' % p
      if p32 in e.target:
        new_path = e.target.replace(p32, p)
        showit(show, 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 LIBDIRS:
      if e.path == '/%s' % p:
        e = None
        break


def migrate_package(show, cat, pkg, vdb_pkg, root, dry_run=True, verbose=False):
  """Migrate the contents of |cat|/|pkg|"""
  # For simple packages (like virtuals) there might not be a CONTENTS file.
  contents = os.path.join(vdb_pkg, 'CONTENTS')
  if not os.path.exists(contents):
    if verbose:
      print('SKIP')
    return

  # 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)

      migrate_one_entry(show, cat, pkg, e, root, dry_run=dry_run,
                        verbose=verbose)

      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 needed.
  if not dry_run and modified:
    content = '\n'.join(new_contents)
    if content:
      content += '\n'
    atomic_write(contents, content)


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

  set_symlink_lib_no(root, dry_run=dry_run, verbose=verbose)
  move_libdirs(root, dry_run=dry_run, verbose=verbose)

  show = {
      'cat': False,
      'pkg': False,
  }

  # 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(show, cat, None)

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

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

      migrate_package(show, cat, pkg, vdb_pkg, root, dry_run=dry_run,
                      verbose=verbose)


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 opts.dry_run:
    print('DRY-RUN MODE: no changes will actually be made!')
  elif 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:]))