clean-boot

#!/usr/bin/env python

import os, sys, re, collections as cs


def main(args=None):
	import argparse
	parser = argparse.ArgumentParser(
		description='Tool to cleanup /boot'
			' from older kernels and config/System.map files.')

	parser.add_argument('-d', '--boot-dir',
		metavar='path', default='/boot',
		help='Path to dir where kernels and'
			' config/System.map files are stored (default: %(default)s).')

	parser.add_argument('-f', '--free-space-min',
		type=float, default=20, metavar='percent',
		help='Target min free space amount after cleanup (default: %(default)s).')
	parser.add_argument('-m', '--keep-min',
		type=int, default=3, metavar='count',
		help='Min number of older version to keep (default: %(default)s).')
	parser.add_argument('-o', '--clean-old', action='store_true',
		help='Remove all non-linked ".old" variants, if any.'
			' These are treated same as slightly older versions by default.')
	parser.add_argument('-c', '--keep-configs', action='store_true',
		help='Keep config-* files, removing only vmlinuz-* and System.map-*.')

	parser.add_argument('-r', '--force-remove-count', metavar='count', default=0,
		help='Number of least relevant versions to explicitly remove.'
			' Regardless of this option, --keep-min count is preserved.'
			' "inf" value can be used to remove all but --keep-min count.')

	parser.add_argument('-n', '--dry-run',
		action='store_true', help='Do not actually remove anything.')
	parser.add_argument('--debug', action='store_true', help='Verbose operation mode.')
	opts = parser.parse_args(sys.argv[1:] if args is None else args)

	import logging
	logging.basicConfig(level=logging.DEBUG if opts.debug else logging.WARNING)
	log = logging.getLogger()

	os.chdir(opts.boot_dir)

	path_re = re.compile(r'^(vmlinuz|config|System\.map)-(\d+\.\d+\.\d+)-.*?(\.old)?$')
	path_ver = lambda v,old: tuple(map(int, v.split('.'))) + (int(not bool(old)),)
	format_ver = lambda v: '.'.join(map(str, v))


	## Prepare ordered by-version hashes

	# Hash by-version-tuple (e.g. "3, 12, 20, 1" - last "1" is "0" for ".old")
	# Values are lists of files (vmlinuz, config, System.map) corresponding to each version
	vfiles = cs.defaultdict(list)
	for p in os.listdir('.'):
		m = path_re.search(p)
		if not m: continue
		p_type, v = m.group(1), m.group(2)
		if opts.keep_configs and p_type == 'config': continue
		v = path_ver(v, m.group(3))
		vfiles[v].append(p)
	vfiles = dict(sorted(vfiles.items()))

	# Keep linked version and its .old variant (if any)
	rmq_kept = list()
	if os.path.exists('vmlinuz'):
		p_linked = os.readlink('vmlinuz')
		m = path_re.search(p_linked)
		assert m, p_linked
		v_linked = v = path_ver(m.group(2), m.group(3))
		v_linked_major = v_linked[:2]
		assert v_linked in vfiles, [v, vfiles]
		for old in 0, 1:
			v = v[:3] + (old,)
			if v in vfiles: rmq_kept.append((v, vfiles.pop(v)))

	# Rehash by version_major-version
	vfiles_by_major = dict()
	for v, files in vfiles.items():
		v_major = v[:2]
		if v_major not in vfiles_by_major:
			vfiles_by_major[v_major] = dict()
		vfiles_by_major[v_major][v] = files


	## Build order in which stuff will be removed until there's enough space

	rmq = dict()

	# Prioritize .old variants removal for all but linked version
	if opts.clean_old:
		for v, files in vfiles.items():
			if v[3] == 0: rmq[v] = files

	for v_major, files in vfiles_by_major.items():
		files = list(files.items())[:-1] # keep last patchset of each major ver, if possible
		for v, files in files: rmq[v] = files

	# Remove older majors completely as a last-resort
	for v_major, files in vfiles_by_major.items():
		v, files = list(files.items())[-1]
		rmq[v] = files

	# Keep some number of older versions, if instructed
	if opts.keep_min > 0:
		rmq = list(rmq.items())
		rmq_kept = dict(rmq[-opts.keep_min:] + rmq_kept)
		rmq = dict(rmq[:-opts.keep_min])


	## Removal

	def free():
		df = os.statvfs('.')
		return (df.f_bavail * df.f_bsize / float(df.f_blocks * df.f_bsize)) * 100
	rm_count = opts.force_remove_count
	if rm_count == 'inf': rm_count = 2**64
	else:
		try: rm_count = int(rm_count)
		except: parser.error('Invalid integer value: {}'.format(rm_count))

	rmq_set = set() # check to make sure these don't get repeated anywhere
	if rmq_kept:
		log.debug( 'Preserved versions (linked version,'
			' its ".old" variant, --keep-min): %s', len(rmq_kept) )
		for v, files in rmq_kept.items():
			for p in files:
				log.debug(' - %s - %s', format_ver(v), p)
				rmq_set.add(p)

	df, rmq_len_start = free(), len(rmq)
	for v, files in list(rmq.items()):
		df = free()
		if rm_count > 0: rm_count -= 1 # avoids df check
		elif df >= opts.free_space_min: break
		log.debug('Removing files for version (df: %.1f%%): %s', df, format_ver(v))
		for p in files:
			log.debug(' - %s', p)
			assert p not in rmq_set, p
			if not opts.dry_run: os.unlink(p)
		del rmq[v]

	rmq_len = len(rmq)
	rmq_len_left = rmq_len + len(rmq_kept)
	log.debug(
		'Finished (df: %.1f%%, versions left: %s, versions removed: %s).',
		df, rmq_len_left, rmq_len_start - rmq_len )

if __name__ == '__main__': sys.exit(main())