#!/usr/bin/python

import os
import sys
import optparse
import errno
import logging
import yumbootstrap.log
import fcntl
import time

import yumbootstrap.yum
import yumbootstrap.suites
import yumbootstrap.fs
import yumbootstrap.sh
from yumbootstrap.exceptions import YBError

#-----------------------------------------------------------------------------

SUITE_DIR = '/etc/yumbootstrap/suites'

#-----------------------------------------------------------------------------

logger = logging.getLogger()
logger.addHandler(yumbootstrap.log.ProgressHandler())
#logger.setLevel(logging.WARNING) # this is default loglevel

#-----------------------------------------------------------------------------

o = optparse.OptionParser(
  usage = '\n  %prog [options] <suite> <target>'
          '\n  %prog [options] --print-config <suite> <target> > yum.conf'
          '\n  %prog [options] --list-suites'
          '\n  %prog [options] <suite|target> --list-scripts'
          '\n  %prog [options] <target> --just-scripts [script script ...]'
          '',
  description = 'Install Yum-based distribution in a chroot environment.',
)

#-----------------------------------------------------------

expected_nargs = {
  'install':      (lambda n: n == 2),
  'yum.conf':     (lambda n: n == 2),
  'list_suites':  (lambda n: n == 0),
  'list_scripts': (lambda n: n == 1),
  'scripts':      (lambda n: n >= 1),
  'uninstall':    (lambda n: n == 2),
  #'download':     [?],
  #'second_stage': [?],
  #'tarball':      [?],
}

o.set_defaults(
  action = 'install',
  include = [],
  exclude = [],
  groups = [],
  repositories = {},
)

def add_pkg_list(option, opt, value, parser, attr):
  getattr(parser.values, attr).extend(value.split(','))

def add_kv_list(option, opt, value, parser, attr):
  if '=' not in value:
    raise optparse.OptionValueError('"%s" is not in NAME=VALUE format' % value)
  (k,v) = value.split('=', 1)
  getattr(parser.values, attr)[k] = v

#-----------------------------------------------------------

o.add_option(
  '--print-config',
  action = 'store_const', dest = 'action', const = 'yum.conf',
  help = 'print Yum configuration which will be used by yumbootstrap',
)
o.add_option(
  '--uninstall',
  action = 'store_const', dest = 'action', const = 'uninstall',
  help = 'remove chrooted environment with bindings if exists',
)
o.add_option(
  '--list-suites',
  action = 'store_const', dest = 'action', const = 'list_suites',
  help = 'list available suites and exit',
)
o.add_option(
  '--list-scripts',
  action = 'store_const', dest = 'action', const = 'list_scripts',
  help = 'list scripts in a suite (or in a target created with --no-scripts)'
         ' and exit',
)
o.add_option(
  '--just-scripts',
  action = 'store_const', dest = 'action', const = 'scripts',
  help = 'list available suites and exit',
)

#-----------------------------------------------------------

o.add_option(
  '--verbose',
  action = 'store_true', default = False,
  help = 'be verbose about operations',
)
o.add_option(
  '--noninteractive',
  action = 'store_false', dest = 'interactive', default = True,
  help = 'run in non-interactive mode (e.g. no progress bars)',
)
#o.add_option(
#  '--arch', # TODO
#  action = 'store', default = os.uname()[4],
#  help = 'specify target architecture',
#)
#o.add_option(
#  '--proxy', # TODO
#  action = 'store', type = 'string', default = None,
#  help = 'specify a proxy to use when fetching files',
#  metavar = 'HOST:PORT',
#)
o.add_option(
  '--skip-script',
  action = 'append', dest = 'skip', default = [],
  help = 'skip this post-install script(s) (may be specified multiple times;'
         ' see also --list-scripts)',
  metavar = 'NAME',
)
o.add_option(
  '--no-scripts',
  action = 'store_false', dest = 'run_scripts', default = True,
  help = "don't run any post-install scripts after installation",
)
o.add_option(
  '--suite-dir',
  action = 'store', dest = 'suite_dir', type = 'string', default = SUITE_DIR,
  help = 'specify a directory with suite definitions',
  metavar = 'PATH',
)
o.add_option(
  '--include',
  action = 'callback', type = 'string',
  callback = add_pkg_list, callback_args = ('include',),
  help = 'include these packages (comma separated list; may be specified'
         ' multiple times)',
  metavar = 'RPMS',
)
o.add_option(
  '--exclude',
  action = 'callback', type = 'string',
  callback = add_pkg_list, callback_args = ('exclude',),
  help = 'exclude these packages (comma separated list; may be specified'
         ' multiple times)',
  metavar = 'RPMS',
)
o.add_option(
  '--groups',
  action = 'callback', type = 'string',
  callback = add_pkg_list, callback_args = ('groups',),
  help = 'install these package groups (comma separated list; may be specified'
         ' multiple times)',
)
#o.add_option(
#  '--no-default-rpms', # TODO
#  action = 'store_false', dest = 'install_default_rpms', default = True,
#  help = "don't install default RPMs set (useful for splitting installation"
#         " into several parts; see also --skip-fix-rpmdb and --skip-cleanup"
#         " options)",
#)
o.add_option(
  '--gpg-key',
  action = 'append', dest = 'gpg_keys', default = [],
  help = 'add GPG key as a trusted RPM signing key (may be specified'
         ' multiple times)',
  metavar = 'KEYFILE',
)
o.add_option(
  '--repo',
  action = 'callback', type = 'string',
  callback = add_kv_list, callback_args = ('repositories',),
  help = 'use this Yum repository (may be specified multiple times)',
  metavar = 'NAME=URL',
)
#o.add_option(
#  '--download-only', # TODO
#  action = 'store_const', dest = 'action', const = 'download',
#  help = "download RPMs only, don't install them",
#)
#o.add_option(
#  '--foreign', # TODO
#  action = 'store_true', dest = 'no_scripts', default = False,
#  help = "don't run post-install scripts from RPM (mainly useful for"
#         " non-matching architecture in --arch option)",
#)
#o.add_option(
#  '--second-stage', # TODO
#  action = 'store_const', dest = 'action', const = 'second_stage',
#  help = "finalize the installation started with --foreign option",
#)
#o.add_option(
#  '--make-tarball', # TODO
#  action = 'store_const', dest = 'action', const = 'tarball',
#  help = "make a tarball with RPMs instead of installing them",
#)
#o.add_option(
#  '--unpack-tarball', # TODO
#  action = 'store', dest = 'tarball', default = None,
#  help = "use RPMs from a tarball created with --make-tarball option",
#)

opts, args = o.parse_args()

if not expected_nargs[opts.action](len(args)):
  o.error("wrong number of arguments")

if opts.verbose:
  logger.setLevel(logging.INFO)

#-----------------------------------------------------------------------------

#-----------------------------------------------------------

# helper function
def write_cached_suite(yumbootstrap_dir, suite_dir, suite_name):
  from yumbootstrap.fs import touch
  # just hope nobody uses directory names with whitespaces (NL, space or
  # similar)
  suite_file_content = "suite_name: %s\nsuite_dir: %s" % \
                       (suite_name, os.path.abspath(suite_dir))
  touch(yumbootstrap_dir, 'suite_location', text = suite_file_content)

#-----------------------------------------------------------

# helper function
def read_cached_suite(target):
  yum_conf = yumbootstrap.yum.YumConfig(chroot = target)
  suite_location_file = os.path.join(yum_conf.root_dir, 'suite_location')
  try:
    suite_location = dict([
      line.strip().split(': ', 1)
      for line in open(suite_location_file).readlines()
    ])
    suite = yumbootstrap.suites.load_suite(
      suite_location['suite_dir'],
      suite_location['suite_name'],
    )
    return suite
  except IOError, e:
    if e.errno == errno.ENOENT:
      raise YBError("\"%s\" is not a prepared target", target, exit = 1)
    else:
      raise YBError('errno=%d: %s', e.errno, e.strerror)

#-----------------------------------------------------------

# helper function
def run_post_install_scripts(suite, skip = [], just = []):
  logger.info("executing post-install scripts")
  from yumbootstrap import sh

  skip = set(skip)
  if len(just) == 0:
    just = set(suite.post_install.names())
  else:
    just = set(just)

  for (script_name, script) in suite.post_install:
    if script_name in skip or script_name not in just:
      logger.info("skipping %s", script_name)
      continue
    logger.info("running %s", script_name)
    os.environ['SCRIPT_NAME'] = script_name
    os.environ['SCRIPT_PATH'] = script[0]
    sh.run(script, env = suite.environment)

#-----------------------------------------------------------

# helper function
def set_scripts_environment(opts, target):
  os.environ["TARGET"] = os.path.abspath(target)
  if opts.verbose:
    os.environ["VERBOSE"] = "true"
  else:
    os.environ["VERBOSE"] = "false"

#-----------------------------------------------------------

def get_lock_file():
  lock = open('/var/run/yumbootstrap.lock', 'w+')
  return lock

def read_lock_global(lock):
  logger.info("Try to get copy lock")
  while True:
    try:
        fcntl.flock(lock, fcntl.LOCK_SH | fcntl.LOCK_NB)
        logger.info("I have got the copy lock")
        break
    except IOError as e:
        # raise on unrelated IOErrors
        if e.errno != errno.EAGAIN:
            raise
        else:
            time.sleep(1)
  return lock

def write_lock_global(lock):
  logger.info("Try to get recreate lock")
  while True:
    try:
        fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
        logger.info("I have got the recreate lock")
        break
    except IOError as e:
        # raise on unrelated IOErrors
        if e.errno != errno.EAGAIN:
            raise
        else:
            time.sleep(1)
  return lock

def write_unlock_global(lock):
  fcntl.flock(lock, fcntl.LOCK_UN)

#-----------------------------------------------------------

def make_suite_bindings(suite, target):
  if len(suite.packages.bindings)>0:
    logger.info("Binding items from host system")
    for item in suite.packages.bindings:
      if os.path.exists(item):
        yumbootstrap.fs.mkdir(target + os.sep + item)
        yumbootstrap.sh.run(
          ['mount', '--bind', item, target + os.sep + item],
        )

def make_suite_unbindings(suite, target):
  if len(suite.packages.bindings)>0:
    logger.info("Binding items from host system")
    for item in suite.packages.bindings:
      yumbootstrap.sh.run(
          ['umount', target + os.sep + item],
      )

#-----------------------------------------------------------

def do_uninstall(opts, suite_name, target):
  suite = yumbootstrap.suites.load_suite(opts.suite_dir, suite_name)
  if os.path.exists(target):
    make_suite_unbindings(suite, target)
    yumbootstrap.fs.rm_folder(target)

#-----------------------------------------------------------

def make_incomplete_cache(suite, suite_name):
  if suite.get_cache_path() != "":
    target = suite.get_cache_path() + os.sep + suite_name
    yumbootstrap.fs.touch(target, '.incomplete', text = '# just create')

#-----------------------------------------------------------

def release_incomplete_cache(suite, suite_name):
  if suite.get_cache_path() != "":
    target = suite.get_cache_path() + os.sep + suite_name
    os.remove(target + os.sep + '.incomplete')

#-----------------------------------------------------------

def is_incomplete_cache(suite, suite_name):
  if suite.get_cache_path() != "":
    target = suite.get_cache_path() + os.sep + suite_name
    return os.path.exists(target + os.sep + '.incomplete')

#-----------------------------------------------------------

def do_install(opts, suite_name, target):
  from yumbootstrap.fs import touch
  lock = get_lock_file()
  target_new = target
  suite = yumbootstrap.suites.load_suite(opts.suite_dir, suite_name)

  read_lock_global(lock)
  if suite.get_cache_path() != "" and suite.from_cache(suite_name) == True and is_incomplete_cache(suite, suite_name) == False:
    logger.info("Get suite from cache")
    yumbootstrap.fs.copy_filedir(suite.get_cache_path() + os.sep + suite_name + os.sep + 'cache', target)
    logger.info("operation finished")
    write_unlock_global(lock)
    lock.close()
    make_suite_bindings(suite, target)
    return
  write_unlock_global(lock)
  write_lock_global(lock)
  if suite.get_cache_path() != "":
    if not os.path.exists(suite.get_cache_path()):
      yumbootstrap.fs.mkdir(suite.get_cache_path())
    target = suite.get_cache_path() + os.sep + suite_name
    yumbootstrap.fs.rm_folder(target)
    if not os.path.exists(target + os.sep + 'cache'):
      yumbootstrap.fs.mkdir(target + os.sep + 'cache')
    if not suite.make_cache_info(suite_name):
      target = target_new
    target = target + os.sep + 'cache'

  make_incomplete_cache(suite, suite_name)

  set_scripts_environment(opts, target)

  logger.info("installing %s (release %s) to %s",
              suite.name, suite.release, target)

  logger.info("preparing empty /etc/fstab and /etc/mtab")
  os.umask(022)
  # prepare target directory with an empty /etc/fstab
  touch(target, 'etc/fstab', text = '# empty fstab')
  touch(target, 'etc/mtab')

  if len(opts.repositories) > 0:
    logger.info("using custom repositories: %s",
                ", ".join(sorted(opts.repositories)))
    repositories = opts.repositories
  else:
    logger.info("using built-in repositories")
    repositories = suite.repositories

  yum_conf = yumbootstrap.yum.YumConfig(
    chroot = target,
    repos = repositories,
    env = suite.environment,
    maincfg = suite.maincfg,
  )

  write_cached_suite(yum_conf.root_dir, opts.suite_dir, suite_name)

  # installing works also without adding key, but --nogpgcheck is passed to
  # Yum, so it's generally discouraged
  if len(opts.gpg_keys) > 0 or len(suite.gpg_keys) > 0:
    logger.info("adding GPG keys")
    for keyfile in opts.gpg_keys:
      yum_conf.add_key(keyfile)
    for keyfile in suite.gpg_keys:
      yum_conf.add_key(keyfile)
  else:
    pass # TODO: print warning

  yum = yumbootstrap.yum.Yum(
    chroot = target,
    yum_conf = yum_conf,
    interactive = opts.interactive,
  )

  exclude = suite.packages.exclude + opts.exclude

  # main set of packages (should already include yum and /usr/bin/db_load, so
  # `yum.fix_rpmdb()' works)
  logger.info("installing default packages for %s %s",
              suite.name, suite.release)
  install = suite.packages.install + ['@' + g for g in suite.packages.groups]
  yum.install(install, exclude = exclude)

  # requested additional packages
  if len(opts.include) > 0:
    logger.info("installing additional packages requested from command line")
    yum.install(opts.include, exclude = exclude)
  if len(opts.groups) > 0:
    logger.info("installing additional package groups requested from "
                "command line")
    yum.group_install(opts.groups, exclude = exclude)

  if len(suite.packages.files)>0:
    logger.info("installing items from host system")
    for item in suite.packages.files:
      if not os.path.exists(target + os.pathsep + item):
        yumbootstrap.fs.copy_filedir(item, target + os.sep + item)

  if len(suite.post_install) > 0:
    if opts.run_scripts:
      run_post_install_scripts(suite, skip = opts.skip)
    else:
      logger.info("skipping post-install scripts altogether")

  release_incomplete_cache(suite, suite_name)

  if suite.get_cache_path() != "" and suite.from_cache(suite_name) == True and target != target_new:
    yumbootstrap.fs.copy_filedir(suite.get_cache_path() + os.sep + suite_name + os.sep + 'cache', target_new)
    make_suite_bindings(suite, target_new)

  logger.info("operation finished")
  write_unlock_global(lock)
  lock.close()

#-----------------------------------------------------------

def do_list_suites(opts):
  for suite in yumbootstrap.suites.list_suites(opts.suite_dir):
    print suite

#-----------------------------------------------------------

def do_list_scripts(opts, target):
  if target in yumbootstrap.suites.list_suites(opts.suite_dir):
    suite = yumbootstrap.suites.load_suite(opts.suite_dir, target)
  else:
    suite = read_cached_suite(target)
  for name in suite.post_install.names():
    print name

#-----------------------------------------------------------

def do_scripts(opts, target, *scripts):
  set_scripts_environment(opts, target)
  suite = read_cached_suite(target)
  run_post_install_scripts(suite, skip = opts.skip, just = scripts)

#-----------------------------------------------------------

def do_yum_conf(opts, suite, target):
  logger.setLevel(logging.WARNING) # do not honour --verbose option
  import yumbootstrap.suites
  suite = yumbootstrap.suites.load_suite(opts.suite_dir, suite)

  if len(opts.repositories) > 0:
    repositories = opts.repositories
  else:
    repositories = suite.repositories

  yum_conf = yumbootstrap.yum.YumConfig(chroot = target, repos = repositories)

  if len(opts.gpg_keys) > 0:
    cmd = 'cat %s > %s' % (' '.join(opts.gpg_keys), yum_conf.gpg_keys)
    sys.stdout.write('# remember to put the keys to target directory:\n')
    sys.stdout.write('# %s\n\n' % (cmd))
    for keyfile in opts.gpg_keys:
      yum_conf.add_key(keyfile, pretend = True)

  sys.stdout.write(yum_conf.text())
  sys.stdout.flush()

#-----------------------------------------------------------

#-----------------------------------------------------------------------------

try:

  if opts.action == 'install':
    do_install(opts, *args)
  elif opts.action == 'uninstall':
    do_uninstall(opts, *args)
  elif opts.action == 'list_suites':
    do_list_suites(opts, *args)
  elif opts.action == 'list_scripts':
    do_list_scripts(opts, *args)
  elif opts.action == 'scripts':
    do_scripts(opts, *args)
  elif opts.action == 'yum.conf':
    do_yum_conf(opts, *args)
  else:
    # should never happen
    o.error("unrecognized action: %s" % (opts.action,))
except KeyboardInterrupt:
  pass
except YBError, e:
  print >>sys.stderr, e
  sys.exit(e.code)

#-----------------------------------------------------------------------------
# vim:ft=python
