#!/usr/bin/env python3

# SPDX-FileCopyrightText: 2018 Sebastian Wagner
#
# SPDX-License-Identifier: AGPL-3.0-or-later

# -*- coding: utf-8 -*-
"""
TODO: Fix ordering of dicts
"""
import argparse
import csv
import hashlib
import io
import json
import re
import requests
import sys

from typing import Optional


URL = 'https://raw.githubusercontent.com/certtools/malware_name_mapping/master/mapping.csv'
URL_MALPEDIA = 'https://raw.githubusercontent.com/certtools/malware_name_mapping/master/malpedia.csv'
URL_MISP = 'https://raw.githubusercontent.com/MISP/misp-galaxy/main/clusters/threat-actor.json'

REGEX_FROM_HUMAN = re.compile(r"((?P<res1>[a-z])(?=[A-Z])|"  # "fooBar"
                              r"(?P<res2>.)(\\ )(?=[^\]])|"  # "foo bar" but not "foo[-_ ]?bar"
                              r"(?P<res3>[^\[-])\\-(?=[^-]))")  # "foo-bar" but not "foo[-_ ]?bar"
IDENTIFIER_FROM_HUMAN = re.compile(r"[^a-z0-9]+")


def none_len(arg: Optional[list]):
    if arg:
        return len(arg)
    else:
        return 0


def generate_rule(expression: str, identifier: str, name=None):
    return {"rulename": name if name else "%s-%s" % (identifier,
                                                     hashlib.sha1(expression.encode()).hexdigest()[:10]),
            "if": {"classification.taxonomy": "malicious-code",
                   "malware.name": expression
                   },
            "then": {"classification.identifier": identifier, }}


def generate_regex_from_human(*values):
    """
    generate a regular expression from one or more human-readable names

    Steps:
        * Escapes all strings for regex use
        * replaces all spaces with "[-_ ]?"
        * replaces all changes from lower to upper case with "[-_ ]?"
        * ensures the output is all lower case
    """
    values = [re.escape(value) for value in values]
    newvalues = []
    for value in values:
        if not value:
            continue
        match = REGEX_FROM_HUMAN.search(value)
        i = 0
        while match:
            if i > 10:
                break
            replacement = match.group('res1') or match.group('res2') or match.group('res3')
#            print(i, value, match, match.groups(), replacement)
            value = REGEX_FROM_HUMAN.subn(r"%s[-_ ]?" % replacement,
                                          value, count=1)[0]
            match = REGEX_FROM_HUMAN.search(value)
            i += 1
        newvalues.append(value.lower())

    if not newvalues:
        raise ValueError("No or only empty strings given. Can't generate a regular expression.")

    return "^(%s)$" % "|".join(newvalues)


def download(url: str = URL, add_default=False, params=None, include_malpedia=False,
             include_misp=False, mwnmp_ignore_adware=False):
    download = requests.get(url)
    download.raise_for_status()
    rules = [generate_rule(*line[:2])
             for line in csv.reader(io.StringIO(download.text))
             if not line[0].startswith('#')
             if (not mwnmp_ignore_adware or not any(['adware' == field for field in line]))]
    if include_malpedia:
        download = requests.get(URL_MALPEDIA)
        download.raise_for_status()
        rules.extend([generate_rule(*line[:2])
                      for line in csv.reader(io.StringIO(download.text))
                      if not line[0].startswith('#')])
    if include_misp:
        download = requests.get(URL_MISP)
        download.raise_for_status()
        for actor in download.json()['values']:
            names = [actor["value"]] + actor.get("meta", {}).get("synonyms", [])
            identifier = ("%s-generic"
                          "" % IDENTIFIER_FROM_HUMAN.sub("-",
                                                         actor["value"].lower()))
            rule_name = "misp-threat-actors-%s" % identifier

            rules.append(generate_rule(generate_regex_from_human(*names),
                                       identifier,
                                       name=rule_name))

    if add_default:
        rules.append(generate_rule(".*", add_default, name="default"))

    if params:
        rules.extend((generate_rule(rule[0][0], rule[1][0]) for rule in params))

    return json.dumps(rules, indent=4, separators=(',', ': '))


def create_parser():
    parser = argparse.ArgumentParser(
        prog='download_mapping',
        description='Downloads malware family mapping and converts it to modify syntax.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
You can specify additional rules to be added to the file by using:
-e "^some-expression$" -i "some-identifier"
and multiple ones:
-e "^some-expression$" -i "some-identifier" -e "^other-expression$" -i "other-identifier"
'''
    )

    parser.add_argument('filename', nargs='?',
                        help='The filename to write the converted mapping to. If not given, printed to stdout.')
    parser.add_argument('--url', '-u',
                        default=URL,
                        help='The URL to download the mapping from.')
    parser.add_argument('--add-default', '-d',
                        help='Add a default rule to use the malware name as identifier.',
                        # action='store_true')
                        const='{msg[malware.name]}',
                        nargs='?')
    parser.add_argument('--expression', '-e',
                        nargs=1, action='append',
                        help='Expression for an additional rule.',
                        default=[])
    parser.add_argument('--identifier', '-i',
                        nargs=1, action='append',
                        help='Identifier for an additional rule.',
                        default=[])
    parser.add_argument('-m', '--include-malpedia',
                        default=False, action='store_true',
                        help='Include malpedia data (CC BY-NC-SA 3.0), '
                             'see https://malpedia.caad.fkie.fraunhofer.de/usage/tos from %s'
                             '' % URL_MALPEDIA)
    parser.add_argument('--include-misp',
                        default=False, action='store_true',
                        help='Include MISP galaxy cluster data (CC0 1.0 Universal), '
                             'only threat_actor.json is used'
                             'from %s.' % URL_MISP)
    parser.add_argument('--mwnmp-ignore-adware',
                        action='store_true',
                        help='Ignore adware in the malware name mapping data. '
                             'for rules. Ignores all lines with an explicit '
                             "'adware' field.")
    return parser


def main(args):
    if none_len(args.expression) != none_len(args.identifier):
        print('Additional rules must always consist of an expression and an identifier. '
              'Got %d expressions and %d identifiers.' % (len(args.expression), len(args.identifier)),
              file=sys.stderr)
        sys.exit(1)
    rules = download(url=args.url, add_default=args.add_default,
                     params=zip(args.expression, args.identifier),
                     include_malpedia=args.include_malpedia,
                     include_misp=args.include_misp,
                     mwnmp_ignore_adware=args.mwnmp_ignore_adware,
                     )
    if args.filename:
        try:
            with open(args.filename, 'wt') as output:
                output.write(rules)
        except PermissionError:
            print('Could not write to file %r.' % args.filename, file=sys.stderr)
            sys.exit(1)
    else:
        return rules


if __name__ == '__main__':
    parser = create_parser()
    args = parser.parse_args()
    retval = main(args)
    if retval:
        print(retval)
