mkroot

#!/usr/bin/python3

import apt
import argparse
import glob
import hashlib
import os
import os.path
import platform
import random
import shutil
import stat
import subprocess
import sys
import tarfile
import urllib.request, urllib.parse, urllib.error

# Class that enumerates all files installed by a certain Debian package.
# Can work both recursively and non-recursively.
class FileResolver:
  def __init__(self, exclude_files=None):
    self.apt_cache = apt.Cache()
    self.files = {}
    if exclude_files is None:
      self.exclude_files = ['/usr/share']
    else:
      self.exclude_files = exclude_files[:]

  def package_exists(self, package_name):
    return package_name in self.apt_cache

  def files_for(self, package_names, exclude_packages=None, recursive=True,
      include=None, exclude_files=None):
    result = set()
    if isinstance(package_names, str):
      package_names = [package_names]
    if not exclude_packages:
      exclude_packages = []
    for name in package_names:
      result.update(self.__files_for(name, exclude_packages=exclude_packages,
        recursive=recursive))
    excluded = set()
    if include:
      for f in result:
        if not any(f.startswith(x) for x in include):
          excluded.add(f)
    if exclude_files:
      for f in result:
        if any(f.startswith(x) for x in exclude_files):
          excluded.add(f)
    result.symmetric_difference_update(excluded)
    return result

  def __files_for(self, name, exclude_packages, recursive, identation=1):
    if name in exclude_packages:
      return set()
    if name in self.files:
      return self.files[name]
    pkg = self.apt_cache[name]
    self.files[name] = set()
    for f in pkg.installed_files:
      if not f or any(f.startswith(x) for x in self.exclude_files): continue
      if not os.path.exists(f): continue
      st = os.lstat(f)
      if stat.S_ISDIR(st.st_mode): continue
      self.files[name].add(f)
    if recursive and pkg.installed:
      for or_dep in pkg.installed.get_dependencies('Depends'):
        for dep in or_dep:
          self.files[name].update(self.__files_for(dep.name, exclude_packages,
            recursive, identation + 1))
          break
    return self.files[name]

# Utilities to set up a chroot jail filesystem for minijail.  Files are
# hardlinked instead of copied when possible.
class Chroot:
  def __init__(self, path, mountpoint, link=True):
    self.__link = link
    self.chroot = path
    self.mountpoint = mountpoint
    if not self.mountpoint.endswith('/'):
      self.mountpoint += '/'
    if os.path.exists(path):
      shutil.rmtree(path, True)
    os.makedirs(path)

  def __chroot_path(self, path, relative_to=None):
    root = relative_to
    if root is None:
      root = self.mountpoint
    assert path.startswith(root.rstrip('/')), '%s does not start with %s' % (path, root)
    return self.chroot + path[len(root)-1:]

  def file_exists(self, path, relative_to=None):
    return os.path.lexists(self.__chroot_path(path, relative_to))

  def mkdir(self, path, relative_to=None):
    path = os.path.join(self.chroot, os.path.relpath(self.__chroot_path(path, relative_to),
      self.chroot))
    if not os.path.isdir(path):
        os.makedirs(path)

  def touch(self, path, relative_to=None):
    self.mkdir(os.path.dirname(path), relative_to)
    path = os.path.join(self.chroot, os.path.relpath(self.__chroot_path(path, relative_to),
      self.chroot))
    if not os.path.isfile(path):
        with open(path, 'w') as f:
            pass

  def copyfromhost(self, path, relative_to=None, exclude=[]):
    if '*' in path:
      for p in glob.glob(path):
        self.copyfromhost(p, relative_to, exclude)
      return
    try:
      self.mkdir(os.path.dirname(path), relative_to)
    except AssertionError:
      # The roots might try to create their parent directory,
      # which will be outside the root.
      pass
    if os.path.isdir(os.path.realpath(path)):
      for root, dirs, files in os.walk(path):
        if any([root.startswith(e) for e in exclude]): continue
        self.mkdir(root, relative_to)
        for f in files:
          filepath = os.path.join(root, f)
          if os.path.islink(filepath):
            target = os.readlink(filepath)
            abspath = os.path.realpath(os.path.join(os.path.dirname(filepath), target))
            if abspath.startswith(self.mountpoint):
              self.symlink(filepath, target)
            else:
              try:
                self.install(filepath, abspath, relative_to)
              except OSError as e:
                print(e, abspath, filepath)
          else:
            self.install(filepath, os.path.realpath(filepath), relative_to)
        for d in dirs:
          dirpath = os.path.join(root, d)
          if os.path.islink(dirpath):
            target = os.readlink(dirpath)
            abspath = os.path.realpath(os.path.join(os.path.dirname(dirpath), target))
            if abspath.startswith(self.mountpoint):
              self.symlink(dirpath, target)
            else:
              shutil.copytree(abspath, self.__chroot_path(dirpath, relative_to))
    else:
      self.install(path, os.path.realpath(path), relative_to)

  def install(self, path, source, relative_to=None):
    try:
      self.mkdir(os.path.dirname(path), relative_to)
      while os.path.islink(source):
        source = os.path.join(os.path.dirname(source),
          os.readlink(source))
      if os.path.isdir(source):
        self.mkdir(path, relative_to)
      elif self.__link:
        os.link(source, self.__chroot_path(path, relative_to))
      else:
        shutil.copy(source, self.__chroot_path(path, relative_to))
    except Exception as e:
      print(path, source, file=sys.stderr)
      raise e

  def symlink(self, path, destination):
    self.mkdir(os.path.dirname(path))
    try:
      os.symlink(destination, self.__chroot_path(path))
    except Exception as e:
      print(path, destination, file=sys.stderr)
      raise e

  def mknod(self, path, mode=0o600, device=0, *, dir_fd=None):
    self.mkdir(os.path.dirname(path))
    try:
      os.mknod(self.__chroot_path(path), mode=mode, device=device, dir_fd=dir_fd)
      os.chmod(self.__chroot_path(path), mode&0o777)
    except Exception as e:
      print(path, destination, file=sys.stderr)
      raise e

  def write(self, path, contents):
    self.mkdir(os.path.dirname(path))
    with open(self.__chroot_path(path), 'wb') as f:
      f.write(contents)
    os.utime(self.__chroot_path(path), (0, 0))

  def extract(self, archive, sha1, skipprefix, url, exclude=[]):
    if not os.path.exists(archive):
      urllib.request.urlretrieve(url, archive)
    with open(archive, 'rb') as f:
      assert(sha1 == hashlib.sha1(f.read()).hexdigest())
    with tarfile.open(archive) as tar:
      for member in tar.getmembers():
        path = os.path.normpath(member.name)
        if not path.startswith(skipprefix): continue
        path = self.mountpoint + path[len(skipprefix):]
        if any([path.startswith(e) for e in exclude]): continue
        if member.issym():
          self.symlink(path, member.linkname)
        elif member.isfile():
          self.mkdir(os.path.dirname(path))
          with open(self.__chroot_path(path), 'w') as dst:
            shutil.copyfileobj(tar.extractfile(member), dst, member.size)
          os.chmod(self.__chroot_path(path), member.mode)

  # __enter__ and __exit__ are just provided to support with for clarity of code.
  def __enter__(self):
    return self

  def __exit__(self, type, value, traceback):
    pass

def main(args):
  resolver = FileResolver()

  ROOT_PACKAGES = ['libc6', 'libcap2', 'libexpat1','libtinfo5', 'zlib1g',
      'libffi6', 'libgmp10', 'libstdc++6', 'libreadline6', 'libssl1.0.0',
      'libyaml-0-2', 'libgdbm3', 'ca-certificates', 'lua5.2', 'libicu55',
      'libuuid1', 'liblzma5', 'libunwind8']
  OPTIONAL_ROOT_PACKAGES = []
  for package_name in OPTIONAL_ROOT_PACKAGES:
    if resolver.package_exists(package_name):
       ROOT_PACKAGES.append(package_name)

  COMPILER_PACKAGES = ROOT_PACKAGES + ['g++', 'gcc', 'libc6-dev', 'libgmp-dev',
      'libbsd-dev', 'fp-compiler', 'fp-units-fcl', 'libgfortran3', 'libffi-dev',
      'binutils']
  OPTIONAL_COMPILER_PACKAGES = ['libitm1', 'libquadmath0', 'libtsan0']
  for package_name in OPTIONAL_COMPILER_PACKAGES:
    if resolver.package_exists(package_name):
       COMPILER_PACKAGES.append(package_name)

  RUBY_ROOT = '/usr/lib/ruby/'
  HASKELL_ROOT = '/usr/lib/ghc/'
  PYTHON_ROOT = '/usr/lib/python2.7/'
  JAVA_ROOT = '/usr/lib/jvm/'
  DOTNET_ROOT = '/usr/share/dotnet/'

  RUBY_FILES = resolver.files_for('ruby', exclude_packages=COMPILER_PACKAGES)
  PYTHON_FILES = resolver.files_for('python2.7')
  HASKELL_FILES = resolver.files_for('ghc', exclude_packages=COMPILER_PACKAGES)
  JAVA_FILES = resolver.files_for(['openjdk-8-jdk', 'openjdk-8-jre',
      'openjdk-8-jre-headless', 'openjdk-8-jdk-headless'], recursive=False, exclude_files=['/etc'])
  # |resolver| excludes files from /usr/shared. All the dotnet files are in
  # that directory, so we need a new resolver that does not exclude those
  # files.
  raw_resolver = FileResolver(exclude_files=())
  DOTNET_FILES = raw_resolver.files_for(['dotnet-dev-1.0.4', 'dotnet-host',
      'dotnet-hostfxr-1.0.1', 'dotnet-hostfxr-1.1.0',
      'dotnet-sharedframework-microsoft.netcore.app-1.0.5',
      'dotnet-sharedframework-microsoft.netcore.app-1.1.2'], recursive=False)

  def install_common(root):
    if not os.path.exists('/usr/lib/locale/locale-archive'):
      subprocess.check_call(['/usr/sbin/locale-gen', '--purge', 'en_US.UTF-8'])
    root.copyfromhost('/usr/lib/locale/locale-archive')
    root.copyfromhost('/etc/localtime')

    root.write('/etc/passwd',
      b'root:x:0:0:root:/:/bin/false\n'
      b'nobody:x:65534:65534:nobody:/nonexistent:/bin/false')

    # Mountpoints for libraries
    root.mkdir(RUBY_ROOT)
    root.mkdir(PYTHON_ROOT)
    root.mkdir(JAVA_ROOT)
    root.mkdir(HASKELL_ROOT)
    root.mkdir(DOTNET_ROOT)
    root.mkdir('/opt/nodejs/')

    # Other miscellaneous mountpoints
    root.mkdir('/dev/')
    root.mkdir('/proc/')
    root.mkdir('/home/')
    root.mkdir('/sys/')
    root.touch('/dev/stdout')
    root.touch('/dev/stderr')
    root.touch('/dev/stdin')

    # Node.js really wants a source of randomness.
    root.symlink('/dev/urandom', '/opt/nodejs/urandom')

    # Device nodes
    root.mknod('/dev/zero', mode=0o666|stat.S_IFCHR, device=os.makedev(1, 5))

    root.symlink('/usr/bin/java',
            [x for x in JAVA_FILES if x.endswith('bin/java')][0])
    root.symlink('/usr/bin/python', os.path.join(PYTHON_ROOT, 'python'))
    root.symlink('/usr/bin/node', '/opt/nodejs/node')
    root.symlink('/usr/lib/node_modules', '/opt/nodejs/node_modules')
    for f in RUBY_FILES:
        if f.startswith(RUBY_ROOT): continue
        root.symlink(f, os.path.join(RUBY_ROOT, f.replace('/', '_')))

  with Chroot(os.path.join(args.target, 'root'), '/', link=args.link) as root:
    install_common(root)

    for filename in resolver.files_for(ROOT_PACKAGES):
        root.copyfromhost(filename)

    # /tmp is an (optional) mountpoint in the normal chroot, and
    # will be left read-only for the programs that don't need it.
    root.mkdir('/tmp/')

    root.install('/usr/bin/ldwrapper', os.path.join(args.target, 'bin/ldwrapper'))
    root.install('/lib/libminijailpreload.so',
        os.path.join(args.target, 'bin/libminijailpreload.so'))
    root.symlink('/usr/bin/lua', 'lua5.2')

  with Chroot(os.path.join(args.target, 'root-compilers'), '/', link=args.link) as root:
    install_common(root)

    for f in resolver.files_for(COMPILER_PACKAGES):
        if os.path.islink(f):
            root.symlink(f, os.readlink(f))
        else:
            root.copyfromhost(f)

    root.copyfromhost('/usr/bin/fpc')
    root.copyfromhost('/etc/fpc.cfg')
    if os.path.exists('/usr/bin/ppcarm'):
      root.copyfromhost('/usr/bin/ppcarm')
    else:
      root.copyfromhost('/usr/bin/ppcx64')
    root.symlink('/usr/bin/javac',
            [x for x in JAVA_FILES if  x.endswith('bin/javac')][0])
    for f in HASKELL_FILES:
        if f.startswith(HASKELL_ROOT): continue
        if root.file_exists(f): continue
        root.symlink(f, os.path.join(HASKELL_ROOT, f.replace('/', '_')))
    root.symlink('/usr/lib/haskell-packages', os.path.join(HASKELL_ROOT, 'haskell-packages'))
    root.mkdir('/tmp/')

    root.install('/usr/bin/ldwrapper', os.path.join(args.target, 'bin/ldwrapper'))
    root.install('/lib/libminijailpreload.so',
        os.path.join(args.target, 'bin/libminijailpreload.so'))
    root.symlink('/usr/bin/luac', 'luac5.2')

  with Chroot(os.path.join(args.target, 'root-openjdk'), JAVA_ROOT, link=args.link) as root:
    for filename in JAVA_FILES:
        root.copyfromhost(filename)

  with Chroot(os.path.join(args.target, 'root-python'), PYTHON_ROOT, link=args.link) as root:
    root.install(os.path.join(PYTHON_ROOT, 'python'), '/usr/bin/python2.7')
    for filename in PYTHON_FILES:
        if filename.startswith(PYTHON_ROOT):
            root.copyfromhost(filename)

    with urllib.request.urlopen(
            'https://raw.githubusercontent.com/omegaup/karel.js/master/libkarel/libkarel.py') as f:
        root.write('/usr/lib/python2.7/libkarel.py', f.read())

  with Chroot(os.path.join(args.target, 'root-ruby'), RUBY_ROOT, link=args.link) as root:
    root.install(os.path.join(RUBY_ROOT, 'ruby'), '/usr/bin/ruby')
    for filename in RUBY_FILES:
        if filename.startswith(RUBY_ROOT):
            root.copyfromhost(filename)
    for f in RUBY_FILES:
        if f.startswith(RUBY_ROOT): continue
        root.install(os.path.join(RUBY_ROOT, f.replace('/', '_')), f)

  with Chroot(os.path.join(args.target, 'root-hs'), HASKELL_ROOT, link=args.link) as root:
    for filename in HASKELL_FILES:
        if filename.startswith(HASKELL_ROOT):
            root.copyfromhost(filename)
    for f in HASKELL_FILES:
        if f.startswith(HASKELL_ROOT): continue
        root.install(os.path.join(HASKELL_ROOT, f.replace('/', '_')), f)
    root.copyfromhost('/usr/lib/haskell-packages', relative_to='/usr/lib/')

  with Chroot(os.path.join(args.target, 'root-dotnet'), DOTNET_ROOT, link=args.link) as root:
    for filename in DOTNET_FILES:
        if filename.startswith(DOTNET_ROOT):
            root.copyfromhost(filename)
    root.copyfromhost('/usr/share/dotnet/ref', relative_to=DOTNET_ROOT)
    root.copyfromhost('/usr/share/dotnet/Release.rsp', relative_to=DOTNET_ROOT)
    root.write('/usr/share/dotnet/Main.runtimeconfig.json',
               b'{"runtimeOptions": {"framework":{"name":"Microsoft.NETCore.App","version":"1.1.2"}}}')

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Build a chroot environment for minijail')
  parser.add_argument('--target', default='/var/lib/minijail',
                      help='The directory in which the chroot environment will '
                      'be built')
  parser.add_argument('--no-link', dest='link', action='store_false',
                      help='Copy instead of linking files')
  parser.add_argument('command', nargs=argparse.REMAINDER)
  args = parser.parse_args()
  sys.exit(main(args))

# vim: set expandtab:ts=2:sw=2