#!/usr/bin/python3

#	cve-manager : CVE management tool
#	Copyright (C) 2017-2025 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/>.

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

import argparse
import cve_manager.desc        as desc
from time                      import time, localtime, strftime, sleep
from subprocess                import run as run_subprocess
from sys                       import argv
from ax.printer                import Printer
from cve_manager.common        import Init, NewArgParser
from cve_manager.conf          import COMMON_SEC
from cve_manager.intf_download import CmdDownload
from cve_manager.intf_import   import CmdImport
from cve_manager.intf_map      import CmdMap

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Constants

# Names of the modules (without 'cve' prefix)
MANAGER  = 'manager'
BACKUP   = 'backup'
HISTORY  = 'history'
DOWNLOAD = 'download'
IMPORTER = 'import'
MAPPER   = 'map'
ISSUES   = 'issues'
MONITOR  = 'monitor'
MODULES_P1 = [BACKUP, HISTORY, DOWNLOAD, IMPORTER]
MODULES_P2 = [MAPPER, ISSUES]
MODULES = MODULES_P1 + MODULES_P2

# Paths of the modules
PATHS = {
	module: ('cpe-' if module == MAPPER else 'cve-') + module
	for module in MODULES
	}
DEBUG_PATHS = {
	module: './' + ('cve-import/bin/' if module == IMPORTER else '') + path
	for module, path in PATHS.items()
	}

# Descriptions of the modules
DESCRIPTIONS = {
	MANAGER:  desc.MANAGER,
	BACKUP:   desc.BACKUP,
	HISTORY:  desc.HISTORY,
	DOWNLOAD: desc.DOWNLOAD,
	IMPORTER: desc.IMPORTER,
	MAPPER:   desc.MAPPER,
	ISSUES:   desc.ISSUES,
	MONITOR:  desc.MONITOR,
	}

# Number of seconds before each new attempt to perform a step without errors
PAUSE = 300

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Parsing the arguments

argparser = argparse.ArgumentParser(description=desc.MANAGER)
argparser.add_argument(
	'-a', '--run_all',
	action='store_true',
	help='Run all the modules'
	)
argparser.add_argument(
	'-b', '--beginning_step',
	metavar='MODULE_NAME', type=str, choices=MODULES,
	help='Beginning step (so you could skip some of the early stages)'
	)
argparser.add_argument(
	'-e', '--ending_step',
	metavar='MODULE_NAME', type=str, choices=MODULES,
	help='Ending step (so you could skip some of the final stages)'
	)
argparser.add_argument(
	'-o', '--offline',
	action='store_true',
	help='Do not download vulnerability lists, acls, etc.'
	)
argparser.add_argument(
	'-r', '--retry',
	metavar='N_REPEATED_ATTEMPTS', type=int, default=0,
	help='Number of repeated attempts in case of a failure of executed step'
	)
argparser.add_argument(
	'-l', '--list_modules',
	action='store_true',
	help='List available modules (with descriptions)'
	)
argparser = NewArgParser(ptype='m', base=argparser)
args = argparser.parse_args()

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Initializing the printer and reading the conf file

printer, conf = Init(args)
data_sources, err = conf.SelectDataSources()
if err:
	printer.Err(err)
common_params, = conf.Get([COMMON_SEC])

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Print descriptions of available modules

def ListModules():

	for module in MODULES:
		print(module + ':')
		for line in DESCRIPTIONS[module].splitlines():
			print('\t' + line)

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Get an index of a module with a specified name

def IndexOfTheStep(module_name):

	for i, mn in enumerate(MODULES):
		if module_name == mn:
			return i

	return -1

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# Run a given module

def RunModule(module, params, printer):

	return_code = -1

	if module in MODULES:
		path = DEBUG_PATHS[module] if args.debug else PATHS[module]
		command = [path] + params.strip().split()
		if args.debug:
			command.append('--debug')
		if args.plain:
			command.append('--plain')
		if args.silent:
			command.append('--silent')
		printer.Div()
		printer.LineEnd(f'({strftime("%Y-%m-%d, %H:%M:%S", localtime())}) '
			f'Running "{" ".join(command)}"')
		t0 = time()
		try:
			completed_process = run_subprocess(command)
		except FileNotFoundError:
			printer.Err(f'Can\'t find {path}, perhaps the "--debug" flag is '
				f'{"unintentional" if args.debug else "missing"}')
			exit(1)
		dt = time() - t0
		printer.LineEnd(f'({dt:.2f} sec)'.replace('.', ','))
		if completed_process:
			return_code = completed_process.returncode
	else:
		printer.Err(f'There is no "{module}" module')

	return return_code

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

if __name__ == '__main__':

	if len(argv) < 2:
		argparser.print_help()
		exit(1)

	if args.list_modules:
		ListModules()
		exit(0)

	if not args.run_all and not args.beginning_step and not args.ending_step:
		argparser.error('Missing required flag "-a" or "-b" or "-e"')
		exit(1)

	if args.beginning_step and args.ending_step and \
			(IndexOfTheStep(args.beginning_step) > IndexOfTheStep(args.ending_step)):
		argparser.error(f'The "{args.beginning_step}" step goes after '
			f'the "{args.ending_step}" step')
		exit(1)

	t0 = time()

	# Getting the parameters that would latter be passed to the modules
	bak_storage = common_params.get('nbackups_to_store', 2)
	history_storage = common_params.get('ndays_to_store_history_tables', 15)

	# Forming the dict of params used for running each module
	params = {
		BACKUP   : lambda code=-1 : f'--backup --store {bak_storage}',
		HISTORY  : lambda code=-1 : f'--store {history_storage}',
		DOWNLOAD : lambda code=-1 : CmdDownload(code),
		IMPORTER : lambda code=-1 : CmdImport(code),
		MAPPER   : lambda data_source=None :
			lambda data_source=data_source, code=-1 : \
				(f'--data_sources {data_source} ' if data_source else '') + \
				CmdMap(code),
		ISSUES   : lambda data_source=None : \
			lambda data_source=data_source, code=-1 : \
				(f'--types {data_source} ' if data_source else '') + \
				'--branches all --prepare',
		}

	# Forming the steps that imply the running the 1st group of modules
	steps = [(m, params[m]) for m in MODULES_P1
		if not args.offline or m != DOWNLOAD]

	# Adding the steps that imply the running the 2nd group of modules
	for data_source in data_sources:
		steps += [(m, params[m](data_source=data_source)) for m in MODULES_P2]

	# Removing the steps that asked to be skipped
	if not args.run_all:
		steps_specified = []
		ib = IndexOfTheStep(args.beginning_step)
		ie = IndexOfTheStep(args.ending_step)
		for step, params_for_this_step in steps:
			i = IndexOfTheStep(step)
			if (i >= ib or ib < 0) and (i <= ie or ie < 0):
				steps_specified.append((step, params_for_this_step))
		steps = steps_specified

	# Executing the steps
	printer = Printer(silent=args.silent, plain=args.plain)
	for step, params_for_this_step in steps:
		return_code = -1
		count = 0
		while True:
			_params_for_this_step = params_for_this_step(code=return_code)
			return_code = RunModule(step, _params_for_this_step, printer)
			if return_code == 0:
				break
			if count == args.retry:
				exit(1)
			count += 1
			printer.LineEnd(f'Retrying in {PAUSE} seconds...')
			sleep(PAUSE)

	dt = time() - t0
	printer.LineEnd(f'\n(Total time {dt:.2f} sec)'.replace('.', ','))

	exit(0)
