#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
TODO: Fix ordering of dicts
"""
import argparse
import csv
import hashlib
import io
import json
import requests
import sys


URL = 'https://raw.githubusercontent.com/certtools/malware_name_mapping/master/mapping.csv'


def main(url: str=URL, add_default=False, match_family=False):
    download = requests.get(url)
    download.raise_for_status()
    rules = [{"rulename": "%s-%s" % (line[1], hashlib.sha1(line[0].encode()).hexdigest()[:10]),
              "if": {"classification.taxonomy": "malicious code",
                     "malware.name": line[0]
                     },
              "then": {"classification.identifier": line[1], }}
             for line in csv.reader(io.StringIO(download.text))]
    if match_family:
        rules += [{"rulename": line[1],
                   "if": {"malware.name": "^%s$" % line[1]},
                   "then": {"classification.identifier": line[1], }}
                  for line in csv.reader(io.StringIO(download.text))]

    if add_default:
        rules.append(
            {"rulename": "default",
             "if": {"malware.name": ".*",
                    "classification.identifier": ""},
             "then": {"classification.identifier": "{msg[malware.name]}", }})
    return json.dumps(rules, indent=4, separators=(',', ': '))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        prog='download_mapping',
        description='Downloads malware family mapping and converts it to modify syntax.',
    )

    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.',
                        const=True, action='store_const')
    parser.add_argument('--match-family', '-f',
                        help='Add a rules that also match the family name themselves.',
                        const=True, action='store_const')
    args = parser.parse_args()

    rules = main(url=args.url, add_default=args.add_default, match_family=args.match_family)
    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)
            exit(1)
    else:
        print(rules)
