#!/usr/bin/python3

#	cve-manager : CVE management tool
#	Copyright (C) 2017-2022 Alexey Appolonov
#
#	This program is free software: you can redistribute it and/or modify
#	it under the terms of the GNU General Public License as published by
#	the Free Software Foundation, either version 3 of the License, or
#	(at your option) any later version.
#
#	This program is distributed in the hope that it will be useful,
#	but WITHOUT ANY WARRANTY; without even the implied warranty of
#	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#	GNU General Public License for more details.
#
#	You should have received a copy of the GNU General Public License
#	along with this program.  If not, see <http://www.gnu.org/licenses/>.

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

from os                  import path
from collections         import defaultdict
from cve_manager.common  import GetDef
from cve_manager.defines import NVD_DATA_SRC
from cpe_map.defines     import M_PRESCRIBED
from cpe_map.init        import Init
from cpe_map.common      import NewArgParser
from cpe_map.control     import TerminateIfAbsent, TerminateOrSkip

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Form a dict of mapping prescriptions

def GetPresc(target, conf_dir):

	presc = defaultdict(set)
	warnings = []

	FILE_NAME = f'{"cpe" if target == NVD_DATA_SRC else target}-mapping.csv'
	FILE_PATH = path.join(conf_dir, FILE_NAME)

	# If file exists
	if path.isfile(FILE_PATH):
		got_header = False
		line_n = 0
		# Parsing a file line by line
		with open(FILE_PATH, 'r') as f:
			while True:
				line = f.readline()
				line_n += 1
				if not line:
					break
				# Omitting a header
				if not got_header:
					got_header = True
					continue
				# Skipping empty lines
				line = line.strip()
				if not line:
					continue
				fields = [field.strip() for field in line.split(',')]
				if len(fields) < 2:
					warn = f'Wrong format of "{FILE_PATH}" file ' \
						f'in line {line_n}{f": {line}" if line else ""}'
					warnings.append(warn)
					continue
				package, products = fields[0], fields[1]
				if presc.get(package):
					warn = f'Reassignment of {package} in line {line_n} ' \
						f'of "{FILE_PATH}" file'
					warnings.append(warn)
				presc[package] |= set(products.split())

	return presc, warnings

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Get matches between our packages, NVD/FSTEC names and prescripted pairs

def GetMatches(our_names, their_names, presc):

	matches = defaultdict(set)
	warnings = []

	for package, products in presc.items():
		for product in products:
			if package in our_names and product in their_names:
				matches[package].add(product)
			else:
				# Notifying about unused prescriptions
				warnings.append(f'{package} >> {product}')

	matches = {package: '  '.join(products)
		for package, products in matches.items()}

	return matches, warnings

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

if __name__ == '__main__':

	# Parsing the args and getting the helper objects
	args = NewArgParser(ptype='p').parse_args()
	conf, speaker, mediator = Init(M_PRESCRIBED, args)

	# Forming a dict of prescribed mapping pairs
	all_presc = {}
	conf_dir = GetDef('CONF_DIR', args.debug)
	for data_source in args.data_sources:
		speaker.Say(f'Getting {data_source.upper()} prescriptions')
		all_presc[data_source], warnings = GetPresc(data_source, conf_dir)
		speaker.Status(extra_warn=warnings)
	if not all_presc or all([not presc for presc in all_presc]):
		exit(0)

	# Getting package names
	selection = {package for presc in all_presc.values()
		for package in presc.keys()
		if not args.packages or package in args.packages}
	if not selection:
		msg = 'There are no prescriptions'
		if args.packages:
			msg += ' for a given selection'
		speaker.Say(f'{msg}, nothing to do.', endl=True)
		exit(0)
	our_names = mediator.QueryPackages(selection, args.names)
	TerminateIfAbsent(our_names)

	# For every type of selected data source (nvd, fstec)
	for data_source, presc in all_presc.items():

		# Getting product names
		their_names = mediator.QueryProducts(data_source,
			{product for products in presc.values() for product in products})
		if TerminateOrSkip(their_names):
			continue

		speaker.Op(data_source, our_names, their_names)

		# Getting matches and a list of messages about unused prescriptions
		matches, warnings = GetMatches(our_names, their_names, presc)
		if warnings:
			speaker.Status(warn='Unused prescriptions left', extra=warnings)
		else:
			speaker.Status(f'All {len(presc)} prescriptions been used')

		if args.noupdate:
			speaker.Matches(matches)
			continue

		# Updating the table with the results
		if not mediator.SendMatches(data_source, matches):
			exit(1)

	exit(0)
