path: root/bin/emaint

#!/usr/bin/python -O

import sys, os, time, signal
from optparse import OptionParser, OptionValueError
if not hasattr(__builtins__, "set"):
	from sets import Set as set
import re
try:
	import portage
except ImportError:
	from os import path as osp
	sys.path.insert(0, osp.join(osp.dirname(osp.dirname(osp.realpath(__file__))), "pym"))
	import portage

import portage.const, portage.exception, portage.output
class WorldHandler(object):

	def name():
		return "world"
	name = staticmethod(name)

	def __init__(self):
		self.invalid = []
		self.not_installed = []
		self.invalid_category = []
		self.okay = []
		self.world_file = os.path.join("/", portage.const.WORLD_FILE)
		self.found = os.access(self.world_file, os.R_OK)

	def _check_world(self, onProgress):
		categories = set(portage.settings.categories)
		myroot = portage.settings["ROOT"]
		vardb = portage.db[myroot]["vartree"].dbapi

		world_atoms = open(self.world_file).read().split()
		maxval = len(world_atoms)
		if onProgress:
			onProgress(maxval, 0)
		for i, atom in enumerate(world_atoms):
			if not portage.isvalidatom(atom):
				self.invalid.append(atom)
				if onProgress:
					onProgress(maxval, i+1)
				continue
			okay = True
			if not vardb.match(atom):
				self.not_installed.append(atom)
				okay = False
			if portage.catsplit(atom)[0] not in categories:
				self.invalid_category.append(atom)
				okay = False
			if okay:
				self.okay.append(atom)
			if onProgress:
				onProgress(maxval, i+1)

	def check(self, onProgress=None):
		self._check_world(onProgress)
		errors = []
		if self.found:
			errors += map(lambda x: "'%s' is not a valid atom" % x, self.invalid)
			errors += map(lambda x: "'%s' is not installed" % x, self.not_installed)
			errors += map(lambda x: "'%s' has a category that is not listed in /etc/portage/categories" % x, self.invalid_category)
		else:
			errors.append(self.world_file + " could not be opened for reading")
		return errors

	def fix(self, onProgress=None):
		self._check_world(onProgress)
		errors = []
		try:
			portage.write_atomic(self.world_file, "\n".join(self.okay))
		except portage.exception.PortageException:
			errors.append(self.world_file + " could not be opened for writing")
		return errors

class BinhostHandler(object):

	def name():
		return "binhost"
	name = staticmethod(name)

	def __init__(self):
		myroot = portage.settings["ROOT"]
		self._bintree = portage.db[myroot]["bintree"]
		self._bintree.populate()
		self._pkgindex_file = os.path.join(self._bintree.pkgdir, "Packages")
		from portage import getbinpkg
		self._pkgindex = getbinpkg.PackageIndex()
		f = open(self._pkgindex_file, 'r')
		try:
			self._pkgindex.read(f)
		finally:
			f.close()

	def check(self, onProgress=None):
		missing = []
		cpv_all = self._bintree.dbapi.cpv_all()
		cpv_all.sort()
		maxval = len(cpv_all)
		if onProgress:
			onProgress(maxval, 0)
		pkgindex = self._pkgindex
		missing = []
		for i, cpv in enumerate(cpv_all):
			d = pkgindex.packages.get(cpv)
			if not d or "MD5" not in d:
				missing.append(cpv)
			if onProgress:
				onProgress(maxval, i+1)
		errors = ["'%s' is not in Packages" % cpv for cpv in missing]
		stale = set(pkgindex.packages).difference(cpv_all)
		for cpv in stale:
			errors.append("'%s' is not in the repository" % cpv)
		return errors

	def fix(self, onProgress=None):
		bintree = self._bintree
		cpv_all = self._bintree.dbapi.cpv_all()
		cpv_all.sort()
		missing = []
		maxval = len(cpv_all)
		if onProgress:
			onProgress(maxval, 0)
		pkgindex = self._pkgindex
		missing = []
		for i, cpv in enumerate(cpv_all):
			d = pkgindex.packages.get(cpv)
			if not d or "MD5" not in d:
				bintree.inject(cpv)
			if onProgress:
				onProgress(maxval, i+1)
		stale = set(pkgindex.packages).difference(cpv_all)
		if stale:
			from portage import locks
			pkgindex_lock = locks.lockfile(
				self._pkgindex_file, wantnewlockfile=1)
			try:
				from portage import getbinpkg
				del pkgindex
				self._pkgindex = getbinpkg.PackageIndex()
				f = open(self._pkgindex_file, 'r')
				try:
					self._pkgindex.read(f)
				finally:
					f.close()
				from portage.dbapi.bintree import binarytree
				self._bintree = binarytree(bintree.root, bintree.pkgdir,
					settings=bintree.settings)
				del bintree
				portage.db[self._bintree.root]["bintree"] = self._bintree
				self._bintree._populate()
				for cpv in set(self._pkgindex.packages).difference(
					self._bintree.dbapi.cpv_all()):
					del self._pkgindex.packages[cpv]
				from portage.util import atomic_ofstream
				f = atomic_ofstream(self._pkgindex_file)
				try:
					self._pkgindex.write(f)
				finally:
					f.close()
			finally:
				locks.unlockfile(pkgindex_lock)
		return None

class VdbKeyHandler(object):
	def name():
		return "vdbkeys"
	name = staticmethod(name)

	def __init__(self):
		self.list = portage.db["/"]["vartree"].dbapi.cpv_all()
		self.missing = []
		self.keys = ["HOMEPAGE", "SRC_URI", "KEYWORDS", "DESCRIPTION"]
		
		for p in self.list:
			mydir = os.path.join(os.sep, portage.settings["ROOT"], portage.const.VDB_PATH, p)+os.sep
			ismissing = True
			for k in self.keys:
				if os.path.exists(mydir+k):
					ismissing = False
					break
			if ismissing:
				self.missing.append(p)
		
	def check(self):
		return ["%s has missing keys" % x for x in self.missing]
	
	def fix(self):
	
		errors = []
	
		for p in self.missing:
			mydir = os.path.join(os.sep, portage.settings["ROOT"], portage.const.VDB_PATH, p)+os.sep
			if not os.access(mydir+"environment.bz2", os.R_OK):
				errors.append("Can't access %s" % (mydir+"environment.bz2"))
			elif not os.access(mydir, os.W_OK):
				errors.append("Can't create files in %s" % mydir)
			else:
				env = os.popen("bzip2 -dcq "+mydir+"environment.bz2", "r")
				envlines = env.read().split("\n")
				env.close()
				for k in self.keys:
					s = [l for l in envlines if l.startswith(k+"=")]
					if len(s) > 1:
						errors.append("multiple matches for %s found in %senvironment.bz2" % (k, mydir))
					elif len(s) == 0:
						s = ""
					else:
						s = s[0].split("=",1)[1]
						s = s.lstrip("$").strip("\'\"")
						s = re.sub("(\\\\[nrt])+", " ", s)
						s = " ".join(s.split()).strip()
						if s != "":
							try:
								keyfile = open(mydir+os.sep+k, "w")
								keyfile.write(s+"\n")
								keyfile.close()
							except (IOError, OSError), e:
								errors.append("Could not write %s, reason was: %s" % (mydir+k, e))
		
		return errors

class ProgressHandler(object):
	def __init__(self):
		self.curval = 0
		self.maxval = 0
		self.last_update = 0
		self.min_display_latency = 0.2

	def onProgress(self, maxval, curval):
		self.maxval = maxval
		self.curval = curval
		cur_time = time.time()
		if cur_time - self.last_update >= self.min_display_latency:
			self.last_update = cur_time
			self.display()

	def display(self):
		raise NotImplementedError(self)

def emaint_main(myargv):

	# TODO: Create a system that allows external modules to be added without
	#       the need for hard coding.
	modules = {
		"world" : WorldHandler,
		"binhost":BinhostHandler
	}

	module_names = modules.keys()
	module_names.sort()
	module_names.insert(0, "all")

	def exclusive(option, *args, **kw):
		var = kw.get("var", None)
		if var is None:
			raise ValueError("var not specified to exclusive()")
		if getattr(parser, var, ""):
			raise OptionValueError("%s and %s are exclusive options" % (getattr(parser, var), option))
		setattr(parser, var, str(option))


	usage = "usage: emaint [options] " + " | ".join(module_names)

	usage+= "\n\nCurrently emaint can only check and fix problems with one's world\n"
	usage+= "file.  Future versions will integrate other portage check-and-fix\n"
	usage+= "tools and provide a single interface to system health checks."


	parser = OptionParser(usage=usage)
	parser.add_option("-c", "--check", help="check for problems",
		action="callback", callback=exclusive, callback_kwargs={"var":"action"})
	parser.add_option("-f", "--fix", help="attempt to fix problems",
		action="callback", callback=exclusive, callback_kwargs={"var":"action"})
	parser.action = None


	(options, args) = parser.parse_args(args=myargv)
	if len(args) != 1:
		parser.error("Incorrect number of arguments")
	if args[0] not in module_names:
		parser.error("%s target is not a known target" % args[0])

	if parser.action:
		action = parser.action
	else:
		print "Defaulting to --check"
		action = "-c/--check"

	if args[0] == "all":
		tasks = modules.values()
	else:
		tasks = [modules[args[0]]]


	if action == "-c/--check":
		status = "Checking %s for problems"
		func = "check"
	else:
		status = "Attempting to fix %s"
		func = "fix"

	isatty = sys.stdout.isatty()
	for task in tasks:
		print status % task.name()
		inst = task()
		onProgress = None
		if isatty:
			progressBar = portage.output.TermProgressBar()
			progressHandler = ProgressHandler()
			onProgress = progressHandler.onProgress
			def display():
				progressBar.set(progressHandler.curval, progressHandler.maxval)
			progressHandler.display = display
			def sigwinch_handler(signum, frame):
				lines, progressBar.term_columns = \
					portage.output.get_term_size()
			signal.signal(signal.SIGWINCH, sigwinch_handler)
		result = getattr(inst, func)(onProgress=onProgress)
		if isatty:
			# make sure the final progress is displayed
			progressHandler.display()
			print
			signal.signal(signal.SIGWINCH, signal.SIG_DFL)
		if result:
			print
			print "\n".join(result)
			print "\n"

	print "Finished"

if __name__ == "__main__":
	emaint_main(sys.argv[1:])