purge-distfiles

#!/usr/bin/python

# On each client:
# Look through all the environment.bz2 files under /var/db/pkg/*/*
#  - Find line: declare -x A="<space delimited list>"
#  - Add all the files in this list to a set of in-use files
# Save the in-use set under distfiles, in a file named for the host
#
# Then, on one PC
# - Load all the client file lists into a set, REQUIRED
# - Glob all the filenames in distfiles into a set, PRESENT
# - Find filenames in PRESENT but not in REQUIRED, into a set UNWANTED
# - Delete all the UNWANTED files.

import argparse
import bz2
import glob
import os
import os.path
import platform
import re
import stat
import sys

MAKE_CONF = '/etc/portage/make.conf'
PORTAGE_DATABASE_DIR = '/var/db/pkg'
CLEANUP_FILE = '/root/clean-distfiles'

# This is subdirectory in distfiles where this app keeps it's files
# (All the clients can see distfiles)
MANAGE_DIR = '.purgeman'


class PurgeException(Exception):
  pass


def getfilesize(f):
  """
  Return file size (bytes)
  """

  return os.stat(f)[stat.ST_SIZE]


def to_sane(b):
  """
  Convert sizes to reasonable units for display
  """

  gb = int(b / (1024*1024*1024))
  if gb >= 5:
    return f'{gb} GB'

  mb = int(b / (1024*1024))
  if mb > 0:
    return f'{mb} MB'

  kb = int(b / 1024)
  return f'{kb} KB'


def find_distdir():
  """
  Find DISTDIR value or give up
  """

  distdir_re = re.compile(r'^DISTDIR="([^"]+)')
  try:
    with open(MAKE_CONF, 'r') as conf:
      for line in conf:
        check = distdir_re.match(line)
        if check:
          distdir = check.group(1)
          if os.path.isdir(distdir):
            return distdir

          print(f'Error: DISTDIR ({distdir}) is not a directory')
          sys.exit()

    print(f'Error: Could not find DISTDIR definition in {MAKE_CONF}')
    sys.exit()

  except OSError:
    print(f'Error: Unable to read portage configuration, {MAKE_CONF}')
    sys.exit()


def find_files():
  """
  Get list of distfiles required by this system
  Save list to <DIST_DIR>/<MANAGE_DIR>/<HOSTNAME>.req
  """

  distdir = find_distdir()
  distfiles = set( glob.glob('*', root_dir=distdir) )

  # Note this value matches an empty string too, some are!
  define_A_re = re.compile(r'declare \-x A="([^"]*)"')

  # Absolute path to each environment.bz2 file
  envfiles = glob.glob(
    os.path.join(PORTAGE_DATABASE_DIR, '*/*/environment.bz2')
  )

  wanted = set()
  for envfile in envfiles:
    with bz2.open(envfile, mode='rt') as env:
      found = False
      for line in env:
        check = define_A_re.match(line)
        if check:
          found = True
          sourcefiles = check.group(1)
          if len(sourcefiles) > 0:
            sf_list = sourcefiles.split(' ')
            wanted.update(sf_list)
          continue
      if not found:
        print(f'Error: package without A, {envfile}')
        sys.exit()

  mylistfilename = os.path.join(distdir, MANAGE_DIR, platform.node() + '.reqs')
  with open(mylistfilename, 'w') as reqs:
    for want in wanted:
      reqs.write(f'{want}\n')
  print(f'Updated requirements: {mylistfilename}')


def purge_files(verbose):
  """
  Get the set of all files in DISTDIR
  Get the set of all files in the client lists
  Use set difference to identify unwanted files
  Write script to remove the unwanted files
  """

  distdir = find_distdir()
  distfiles = set( glob.glob('*', root_dir=distdir) )

  referenced = set()
  reqs = glob.glob(os.path.join(distdir, MANAGE_DIR, '*.reqs'))
  for reqfile in reqs:
    with open(reqfile,'r') as req:
      for line in req:
        referenced.add(line.strip())

  # Now, find all files in distfiles that are not in referenced
  unwanted = distfiles.difference(referenced)

  # For listing purposes, it's nicer to have a a sorted list
  delete_these = list(unwanted)
  delete_these.sort()

  total = 0
  with open(CLEANUP_FILE, 'w') as cleaner:

    cleaner.write('#!/bin/bash\n')
    cleaner.write('\n')
    cleaner.write('# Purges unused entries from distfiles, mulitple PCs\n')
    cleaner.write('# (generated by purge-distfiles)\n')
    cleaner.write('\n')

    for item in delete_these:
      target = os.path.join(distdir, item)
      if os.path.isfile(target):
        size = getfilesize(target)
        total += size
        cmd = f'rm {target} # ({to_sane(size)})'
        if verbose:
          print(cmd)
        cleaner.write(cmd)
        cleaner.write('\n')

    cleaner.write('\n')
    cleaner.write(f'rm {CLEANUP_FILE}\n')

  if total > 0:

    # Make the cleanup file executable
    os.chmod(CLEANUP_FILE, stat.S_IRUSR + stat.S_IWUSR + stat.S_IXUSR)

    print(f'\nTotal size to be removed: {to_sane(total)}')
    print(f'To perform cleanup, run: {CLEANUP_FILE}')

  else:

    os.remove(CLEANUP_FILE)
    print('Nothing to do!')


if __name__ == "__main__":

  parser = argparse.ArgumentParser(description='''
            Purge unused files from a shared distfiles directory.

            Run the command from each PC to generate the list of distfiles
            that it uses. The list files, named after the PC, are saved
            under the distfiles directory.

            Run the command using the --purge option from one of the PCs
            to remove any distfiles not needed by any PC.
            ''')
  parser.add_argument('--purge',
                      action='store_true',
                      help='Purge all unused dist files. When not set, '
                           'generate the in-use distfile list for this PC.')
  parser.add_argument('--verbose', '-v',
                      action='store_true',
                      help='Print the list of files that will be removed.')
  options = parser.parse_args()

  if options.purge:
    purge_files(verbose=options.verbose)
  else:
    find_files()