Commit 40df22a8 authored by Ondřej Borýsek's avatar Ondřej Borýsek
Browse files

Extract parse_arguments

parent 919ac516

scan2report.py

+76 −63
Original line number Diff line number Diff line
@@ -123,9 +123,9 @@ COLORS_CONST = ['D9D9D9', '92D050', 'FFFF00', 'FFC000', 'FF0000'] 
debug: bool = False





class Scan2ReportInstance:

class Scan2ReportConfig:

    def __init__(self):

        self.hostnames: dict = {}

        self.hostnames: dict = {}  # todo: remove hostnames from config?

        self.lang: str = 'cs'

        self.ignore: list = []

        self.ignorePluginOutput: bool = False
@@ -133,14 +133,16 @@ class Scan2ReportInstance: 
        self.minSeverity: int = 0

        self.outfile: str = 'out.docx'

        self.outfileFormat: str = 'docx'

        self.pluginStats: dict = {}

        self.pluginStats: dict = {}  # todo: remove pluginStats from config?

        self.pluginStatsOnly: bool = False

        self.profile: str = 'default'

        self.rootDir: str = os.path.dirname(os.path.realpath(__file__))

        self.resolveRedirects: bool = False

        self.skipServices: bool = False

        self.skipTcpWrapped: bool = False

        self.supportedOutfileFormats: List[str] = ['docx', 'csv', 'json']

        self.templateDocDefault: str = 'templates/tns.docx'

        self.templateDoc: str = self.templateDocDefault



    def usage(self):

        print('Usage:', sys.argv[0], '[OPTIONS] .. <FILE> ..\n',
@@ -160,6 +162,56 @@ class Scan2ReportInstance: 
              '       -t, --template <file>: use specified template document, default:', self.templateDocDefault, '\n',

              '       -W, --skip-tcp-wrapped: ignore open ports detected as TCP wrapped, default:', self.skipTcpWrapped, '\n')



    def parse_arguments(self, argv: List[str]):

        # parse arguments

        try:

            opts, args = getopt.getopt(argv, 'dhi:I:l:m:no:p:sSt:W', ['debug', 'help', 'ignore=', 'include=', 'lang=', 'min-severity=', 'output=', 'profile=', 'resolve-redirects', 'skip-services', 'plugin-stats', 'template', 'skip-tcp-wrapped'])

        except getopt.GetoptError:

            self.usage()

            sys.exit(2)

        for opt, arg in opts:

            if opt in ('-d', '--debug'):

                global debug

                debug = True

            elif opt in ('-h', '--help'):

                self.usage()

                sys.exit()

            elif opt in ('-i', '--ignore'):

                self.ignore = arg.split(',')

            elif opt in ('-I', '--include'):

                self.include = arg.split(',')

            elif opt in ('-l', '--lang'):

                self.lang = arg

            elif opt in ('-m', '--min-severity'):

                self.minSeverity = int(arg)

            elif opt in ('-o', '--output'):

                self.outfile = arg

            elif opt in ('-p', '--profile'):

                self.profile = arg

            elif opt in ('-r', '--resolve-redirects'):

                self.resolveRedirects = False

            elif opt in ('-s', '--skip-services'):

                self.skipServices = True

            elif opt in ('-S', '--plugin-stats'):

                self.pluginStatsOnly = True

            elif opt in ('-t', '--template'):

                self.templateDoc = os.path.realpath(arg)

                if not (os.path.isfile(self.templateDoc)):

                    print('File "{}" doesn\'t exist or isn\'t readable'.format(self.templateDoc))

                    sys.exit()

            elif opt in ('-W', '--skip-tcp-wrapped'):

                self.skipTcpWrapped = True

        if not self.templateDoc:

            self.templateDoc = self.rootDir + '/' + self.templateDocDefault



        if self.outfile:

            self.outfileFormat = re.sub(r'.*\.', '', self.outfile.lower())

            if self.outfileFormat not in self.supportedOutfileFormats:

                print(f'Unknown output file format: {self.outfileFormat}')

                sys.exit()



        return args





# add Markdown formatted text

def add_md(document, parent, text):
@@ -285,7 +337,7 @@ def lmh2num(char): 


# get finding severity (for sorting)

def pluginsByFindingSeverity(plugin, lang: str):

    rootDir = os.path.dirname(os.path.realpath(__file__))

    rootDir = os.path.dirname(os.path.realpath(__file__))  # todo: take rootDir from s2r



    if os.path.isfile(rootDir + '/plugins/' + lang + '/' + plugin['id'] + '.json'):

        with open(rootDir + '/plugins/' + lang + '/' + plugin['id'] + '.json', encoding="utf8") as jsonfile:
@@ -400,17 +452,16 @@ def produce_findings_json(findings: dict, hostnames): 
    return answer





def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportInstance' = None):

def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportConfig' = None):

    if reinit_global_vars:

        s2r = Scan2ReportInstance()

        s2r = Scan2ReportConfig()



    global debug

    ignore = s2r.ignore

    ignorePluginOutput = s2r.ignorePluginOutput

    include = s2r.include

    args = s2r.parse_arguments(argv)



    hostnames = s2r.hostnames

    lang = s2r.lang

    minSeverity = s2r.minSeverity

    ignore = s2r.ignore

    outfile = s2r.outfile

    outfileFormat = s2r.outfileFormat

    pluginStats = s2r.pluginStats
@@ -418,15 +469,13 @@ def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportInstance' = No 
    profile = s2r.profile

    resolveRedirects = s2r.resolveRedirects

    skipServices = s2r.skipServices

    skipTcpWrapped = s2r.skipTcpWrapped

    templateDocDefault = s2r.templateDocDefault

    supportedOutfileFormats = s2r.supportedOutfileFormats

    rootDir = s2r.rootDir

    templateDoc = s2r.templateDoc



    fidAliases = {}

    findings = {}

    services = {}

    rootDir = os.path.dirname(os.path.realpath(__file__))

    templateDoc = ''



    unknownCounter = 0

    missingPlugins = []

    boldStart = '**'
@@ -440,52 +489,6 @@ def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportInstance' = No 
            print('Error parsing file ' + rootDir + '/plugins/aliases.json:\n' + format(err))

            sys.exit(1)



    # parse arguments

    try:

        opts, args = getopt.getopt(argv, 'dhi:I:l:m:no:p:sSt:W', ['debug', 'help', 'ignore=', 'include=', 'lang=', 'min-severity=', 'output=', 'profile=', 'resolve-redirects', 'skip-services', 'plugin-stats', 'template', 'skip-tcp-wrapped'])

    except getopt.GetoptError:

        s2r.usage()

        sys.exit(2)

    for opt, arg in opts:

        if opt in ('-d', '--debug'):

            debug = True

        elif opt in ('-h', '--help'):

            s2r.usage()

            sys.exit()

        elif opt in ('-i', '--ignore'):

            ignore = arg.split(',')

        elif opt in ('-I', '--include'):

            include = arg.split(',')

        elif opt in ('-l', '--lang'):

            lang = arg

        elif opt in ('-m', '--min-severity'):

            minSeverity = int(arg)

        elif opt in ('-o', '--output'):

            outfile = arg

        elif opt in ('-p', '--profile'):

            profile = arg

        elif opt in ('-r', '--resolve-redirects'):

            resolveRedirects = False

        elif opt in ('-s', '--skip-services'):

            skipServices = True

        elif opt in ('-S', '--plugin-stats'):

            pluginStatsOnly = True

        elif opt in ('-t', '--template'):

            templateDoc = os.path.realpath(arg)

            if not (os.path.isfile(templateDoc)):

                print('File "{}" doesn\'t exist or isn\'t readable'.format(templateDoc))

                sys.exit()

        elif opt in ('-W', '--skip-tcp-wrapped'):

            skipTcpWrapped = True

    if not templateDoc:

        templateDoc = rootDir + '/' + templateDocDefault



    if outfile:

        outfileFormat = re.sub(r'.*\.', '', outfile.lower())

        if outfileFormat not in supportedOutfileFormats:

            print(f'Unknown output file format: {outfileFormat}')

            sys.exit()



    # select profile

    with open(rootDir + '/plugins/profiles.json', encoding="utf8") as jsonfile:

        try:
@@ -511,7 +514,7 @@ def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportInstance' = No 
        if not (os.path.isfile(infile) and os.access(infile, os.R_OK)):

            print('File "{}" doesn\'t exist or isn\'t readable'.format(infile))

    for infile in args:

        unknownCounter = parse_input_file(debug, fidAliases, findings, hostnames, ignore, ignorePluginOutput, include, infile, lang, minSeverity, missingPlugins, pluginStatsOnly, rootDir, services, skipTcpWrapped, unknownCounter, pluginStats)

        unknownCounter = parse_input_file(debug, s2r, fidAliases, findings, infile, missingPlugins, rootDir, services, unknownCounter)



    # plugin stats

    if pluginStatsOnly:
@@ -602,7 +605,17 @@ def main(argv, reinit_global_vars: bool = False, s2r: 'Scan2ReportInstance' = No 
    print('Report written:', outfile)





def parse_input_file(debug, fidAliases, findings, hostnames, ignore, ignorePluginOutput, include, infile, lang, minSeverity, missingPlugins, pluginStatsOnly, rootDir, services, skipTcpWrapped, unknownCounter, plugin_stats) -> int:

def parse_input_file(debug, s2r: Scan2ReportConfig, fidAliases, findings, infile, missingPlugins, rootDir, services, unknownCounter) -> int:

    hostnames = s2r.hostnames

    ignore = s2r.ignore

    ignorePluginOutput = s2r.ignorePluginOutput

    include = s2r.include

    lang = s2r.lang

    minSeverity = s2r.minSeverity

    pluginStatsOnly = s2r.pluginStatsOnly

    skipTcpWrapped = s2r.skipTcpWrapped

    plugin_stats = s2r.pluginStats



    if not pluginStatsOnly:

        print('Parsing file ' + infile)

    infileFormat = 'unknown'