#!/usr/bin/python3 -sP

from sys import argv, exit, stderr, stdout
from os import getenv, path
from platform import node
from rpm2swidtag import Template, Error, rpm, repodata
import argparse, re
from configparser import ConfigParser
from glob import iglob
import io

TEMPLATE_DIR = "/etc/rpm2swidtag"
CONFIG_FILE = "/etc/rpm2swidtag/rpm2swidtag.conf"

def setup_opts():
	parser = argparse.ArgumentParser(description='Generating SWID tags for rpm packages.')
	group1 = parser.add_mutually_exclusive_group()
	group1.add_argument('-a', '--all', dest='all', action='store_true', help='query all packages with glob pattern')
	group1.add_argument('-p', '--package', dest='rpmfile', action='store_true', help='process rpm package file')
	group1.add_argument('--repo', type=str, metavar='DIR', help="create SWID tag information for yum/dnf repo")
	parser.add_argument('--regid', type=str, help=argparse.SUPPRESS)
	parser.add_argument('--software-creator-from', metavar='FILE', help=argparse.SUPPRESS)
	parser.add_argument('--tag-creator', type=str, metavar='SOURCE-FILE, REGID or "REGID NAME"', help="tagCreator Entity attributes")
	parser.add_argument('--software-creator', type=str, metavar='SOURCE-FILE, REGID or "REGID NAME"', help="softwareCreator Entity attributes")
	parser.add_argument('--distributor', type=str, metavar='SOURCE-FILE, REGID or "REGID NAME"', help="distributor Entity attributes")
	parser.add_argument('--sign-pem', metavar='KEYFILE.pem[,CA.pem[...]]', help="PEM files with key and certificates")
	parser.add_argument('--output-dir', type=str, metavar='DIR',
		help="write SWID tags files into regid subdirectory of DIR; or directly into DIR when the path ends with /.")
	group2 = parser.add_mutually_exclusive_group()
	group2.add_argument('--authoritative', action='store_true', help="produce authoritative tag (per NIST.IR.8060) with Payload, rather than Evidence")
	group2.add_argument('--evidence-deviceid', type=str, dest='device', default=node(), help="Evidence/@deviceId string, defaults to hostname")
	parser.add_argument('--primary-only', action='store_true', help="do not generate supplemental tags")
	parser.add_argument('--print-tagid', action='store_true', help="compute and print tagId(s) to standard output")
	parser.add_argument('--preserve-signing-template', action='store_true', help="keep the XML signing template in the output (for subsequent signing)")
	parser.add_argument('--retain-old-md', metavar='N', action='store', type=int, help="preserve N latest copies of *-swidtags.xml.gz when used with --repo")
	other_group = parser.add_argument_group("config options")
	other_group.add_argument('--config', default=CONFIG_FILE, metavar='FILE', help="location of the configuration file")
	parser.add_argument('package', type=str, nargs='*', help='package(s), glob(s) or file name(s)', metavar='...')
	opts = parser.parse_args()

	error = False
	if opts.tag_creator and (opts.tag_creator.startswith("/") or opts.tag_creator.startswith("./")):
		abs_path = path.abspath(opts.tag_creator)
		if not path.isfile(abs_path):
			stderr.write("%s: --tag-creator specifies file %s which does not exists, skipping\n" % (argv[0], opts.tag_creator))
			error = True
	if opts.software_creator and (opts.software_creator.startswith("/") or opts.software_creator.startswith("./")):
		abs_path = path.abspath(opts.software_creator)
		if not path.isfile(abs_path):
			stderr.write("%s: --software-creator specifies file %s which does not exists, skipping\n" % (argv[0], opts.software_creator))
			error = True
	if opts.distributor and (opts.distributor.startswith("/") or opts.distributor.startswith("./")):
		abs_path = path.abspath(opts.distributor)
		if not path.isfile(abs_path):
			stderr.write("%s: --distributor specifies file %s which does not exists, skipping\n" % (argv[0], opts.distributor))
			error = True
	if error:
		exit(1)

	if opts.regid:
		stderr.write("%s: --regid is obsolete, use --tag-creator instead\n" % argv[0])
	if opts.software_creator_from:
		stderr.write("%s: --software-creator-from is obsolete, use --software-creator instead\n" % argv[0])

	return opts

def load_config(opts):
	config = ConfigParser(delimiters = '=', comment_prefixes = '#')
	config['rpm2swidtag'] = { 'config-file': opts.config, 'template-dir': TEMPLATE_DIR }
	config.read(opts.config)
	config.read(sorted(iglob(opts.config + ".d/*.conf")))
	return config

def load_template(config, opts):
	template = None
	try:
		template_dir = path.join(path.dirname(config.get('rpm2swidtag', 'config-file')),
			config.get('rpm2swidtag', 'template-dir'))
		xml_template = getenv('RPM2SWIDTAG_TEMPLATE', template_dir + "/swidtag-template.xml")
		if opts.print_tagid:
			xslt_file = getenv('RPM2SWIDTAG_XSLT', template_dir + "/rpm2swidtag-tagid.xslt")
		else:
			xslt_file = getenv('RPM2SWIDTAG_XSLT', template_dir + "/rpm2swidtag.xslt")
		return Template(xml_template, xslt_file)
	except Error as e:
		stderr.write("%s: %s\n" % (argv[0], e.strerror))
		exit(5)

def write_out_tag(tag, opts, supplemental=False):
	if opts.output_dir and not opts.print_tagid:
		tag.save_to_directory(opts.output_dir)
	else:
		if opts.print_tagid and supplemental:
			stdout.buffer.write(b"+ ")
		tag.write_output(stdout.buffer)


def process_rpm_with_template(header, template, opts, config, return_xml_list=False):
	params = {}
	if opts.tag_creator:
		if opts.tag_creator.startswith("/") or opts.tag_creator.startswith("./"):
			params['tag-creator-from'] = path.abspath(opts.tag_creator)
		else:
			params['tag-creator-regid'], _, params['tag-creator-name'] = opts.tag_creator.partition(" ")
	if opts.regid:
		params['tag-creator-regid'], params['tag-creator-name'] = opts.regid, None
	if opts.authoritative:
		params['authoritative'] = 'true'
	else:
		params['deviceid'] = opts.device

	signature_key_id = rpm.get_signature_key_id(header)
	component_of = None
	if signature_key_id:
		for s in config.sections():
			if not re.search(r'^link\s', s):
				continue
			if config.get(s, 'key-id') \
				and config.get(s, 'key-id') == signature_key_id:
				creator_from = config.get(s, 'software-creator-file', fallback=None)
				if creator_from:
					creator_from = path.abspath(creator_from)
				if creator_from and path.isfile(creator_from):
					params['software-creator-from'] = creator_from

				distributor_from = config.get(s, 'distributor-file', fallback=None)
				if distributor_from:
					distributor_from = path.abspath(distributor_from)
				if distributor_from and path.isfile(distributor_from):
					params['distributor-from'] = distributor_from

				component_of = config.get(s, 'component-of-file', fallback=None)
				if component_of:
					component_of = path.abspath(component_of)
				if component_of and not path.isfile(component_of):
					component_of = None
				break
	if opts.software_creator:
		if opts.software_creator.startswith("/") or opts.software_creator.startswith("./"):
			params['software-creator-from'] = path.abspath(opts.software_creator)
		else:
			params['software-creator-regid'], _, params['software-creator-name'] = opts.software_creator.partition(" ")
	if opts.software_creator_from:
		params['software-creator-from'] = path.abspath(opts.software_creator_from)

	if opts.distributor:
		if opts.distributor.startswith("/") or opts.distributor.startswith("./"):
			params['distributor-from'] = path.abspath(opts.distributor)
		else:
			params['distributor-regid'], _, params['distributor-name'] = opts.distributor.partition(" ")

	if opts.sign_pem or opts.preserve_signing_template:
		params['preserve-signing-template'] = "true"

	ret = []
	for k in list(params):
		if params[k] is None:
			del params[k]
	tag = template.generate_tag_for_header(header, params=params)
	if opts.sign_pem:
		tag = tag.sign_pem(opts.sign_pem)
	if return_xml_list:
		ret.append(tag)
	else:
		write_out_tag(tag, opts)

	if not opts.primary_only and component_of:
		params['component-of'] = component_of
		tag = template.generate_tag_for_header(header, params=params)
		if opts.sign_pem:
			tag = tag.sign_pem(opts.sign_pem)
		if return_xml_list:
			ret.append(tag)
		else:
			write_out_tag(tag, opts, supplemental=True)

	if return_xml_list:
		return ret

def process_repo(repo, template, opts, config):
	repomd = repodata.Repodata(repo).repomd

	swidtags = repodata.Swidtags(repomd.repo)
	for p in repomd.primary:
		f = path.join(repo, p.href)
		h = rpm.read_from_file(f)
		swidtags.add(p, process_rpm_with_template(h, template, opts, config, return_xml_list=True))
	swidtags.save(retain_old_md=opts.retain_old_md)

def main():
	opts = setup_opts()
	config = load_config(opts)
	template = load_template(config, opts)

	if opts.repo:
		try:
			return process_repo(opts.repo, template, opts, config)
		except Error as e:
			stderr.write("%s: %s\n" % (argv[0], e.strerror))

	if opts.all and not opts.package:
		opts.package = [ '*' ]

	exit_status = 0
	for p in opts.package:
		try:
			if opts.rpmfile:
				l = [ rpm.read_from_file(p) ]
			else:
				# We only assume the use of _RPM2SWIDTAG_RPMDBPATH for testing, really
				l = rpm.read_from_db(p, rpmdb_path=getenv('_RPM2SWIDTAG_RPMDBPATH'),
					glob=opts.all)
		except Error as e:
			stderr.write("%s: %s\n" % (argv[0], e.strerror))
			exit_status = 3
			continue

		found = False
		for h in l:
			if not h["arch"] and not rpm.is_source_package(h):
				continue

			found = True
			try:
				process_rpm_with_template(h, template, opts, config)
				stdout.flush()
			except Error as e:
				if opts.rpmfile:
					msg = "for file [%s]" % p
				else:
					msg = "for package [%s]" % p
				stderr.write("%s: Error generating SWID tag %s: %s\n" % (argv[0], msg, e.strerror))
				exit(6)
			except BrokenPipeError as e:
				stderr.close()
				return exit_status
		if not found:
			stderr.write("%s: No package [%s] found in database\n" % (argv[0], p))
			exit_status = 7

	return exit_status

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