#!/usr/bin/env python3

import argparse
import hashlib
import os
import shutil
import sys

from pycman import config

PACNEW_EXT = '.pacnew'

parser = argparse.ArgumentParser(description='Back up Pacman BACKUP files.')
parser.add_argument('dir', metavar='<dirpath>', help='Output directory.')
parser.add_argument(
  '--config', metavar='<filepath>', default='/etc/pacman.conf',
  help='Pacman configuration file. Default: %(default)s.'
)
parser.add_argument(
  '-r', '--restore', action='store_true',
  help='Restore files in the output directory to the system.',
)
parser.add_argument(
  '--noconfirm', action='store_true',
  help='Do not prompt for confirmation when restoring files.',
)

def confirm(q):
  while True:
    ans = input('{} [Y/n] '.format(q))
    if not ans or ans.lower == 'y':
      return True
    elif ans.lower == 'n':
      return False

def has_changed(p, h):
  d = hashlib.md5()
  with open(p, 'rb') as f:
    chunk = f.read(d.block_size)
    while chunk:
      d.update(chunk)
      chunk = f.read(d.block_size)
  return h != d.hexdigest()

def main(args=None):
  pargs = parser.parse_args(args)
  handle = config.init_with_config(pargs.config)
  db = handle.get_localdb()
  for pkg in db.pkgcache:
    for p, h in pkg.backup:
      ap = os.path.join(handle.root, p)
      op = os.path.join(pargs.dir, p)
      # Do this with a pkg loop to limit restoration to installed packages.
      if pargs.restore:
        if os.path.exists(op):
          try:
            changed = has_changed(ap, h)
            if changed:
              if pargs.noconfirm \
              or confirm('Overwrite {}\n     with {}?'.format(ap, op)):
                shutil.copy2(op, ap)
            else:
              if pargs.noconfirm \
              or confirm('Restore {}\n       to {}?'.format(op, ap)):
                shutil.copy2(op, ap)
                os.rename(ap, ap + PACNEW_EXT)
          except FileNotFoundError:
            if pargs.noconfirm \
            or confirm('Restore {}\n       to {}?'.format(op, ap)):
              shutil.copy2(op, ap)
      else:
        try:
          changed = has_changed(ap, h)
        except FileNotFoundError:
          continue
        if changed:
          try:
            dp = os.path.dirname(op)
            os.makedirs(dp, exist_ok=True)
          # Thrown if permissions differ... makes perfect sense.
          except FileExistsError:
            pass
          # Preserve system directory permissions to protect sensitive system
          # files.
          st = os.stat(os.path.dirname(ap))
          os.chmod(dp, st.st_mode)
          os.chown(dp, st.st_uid, st.st_gid)
          shutil.copy2(ap, op)

if __name__ == '__main__':
  try:
    main()
  except KeyboardInterrupt:
    pass
  except Exception as e:
    sys.exit(str(e))