#!/usr/bin/venv python
# -*- coding: utf-8 -*-

from argparse import ArgumentParser
from config import VT_ATTRIBUTES_MAPPING, PROJECT_NAME
from vt import VT
from dnschecker import DNSChecker as DNS
from emailchecker import EmailChecker
from macchecker import MACChecker 
import ipaddress
from datetime import datetime
from hashing import Hash
from os.path import exists


def checkArguments():
    parser = ArgumentParser(description=PROJECT_NAME)
    parser.add_argument('-c', '--config', help='Config file')
    # For dns command
    parser.add_argument('--dns', help='Get domain name information', action="store_true")
    parser.add_argument('--domain', help='Get domain name information')
    parser.add_argument('--host', help='Get domain name information')
    parser.add_argument('--ip', help='Get IP information')
    # For hash command
    parser.add_argument('--hashfile', help='Hash file', action='store_true')
    parser.add_argument('--scanvt', help='If specified, scan the hash with VirusTotal', action='store_true')
    parser.add_argument('--md5', help='Hash file')
    parser.add_argument('--sha1', help='Hash file')
    parser.add_argument('--sha256', help='Hash file')
    parser.add_argument('--sha384', help='Hash file')
    parser.add_argument('--sha512', help='Hash file')
    parser.add_argument('--hash', help='Get information about the hash')
    # For email command
    parser.add_argument('--email', help='Get email reputation', action='store_true')
    parser.add_argument('--emailrep', help='Get email reputation')
    # For mac command
    parser.add_argument('--mac', help='Get mac information')
    parser.add_argument('--macdb', help='Update database of OUI', action="store_true")

    return parser.parse_args()

def usage():
    print("------------------------------")
    print(f"|           {PROJECT_NAME}           |")
    print("------------------------------\n")
    print("A tool for SOC analyst\n")
    print("Usage: main.py [COMMAND]")
    print("-c PATH, --config PATH\t\tConfig file - mandatory")
    print("--hashfile\t\t\tHash the file and check in VirusTotal")
    print("--hash HASH\t\t\tAnalyse the hash from VirusTotal")
    print("--dns \t\t\t\tGet information regarding the domain with whois and VirusTotal")
    print("--email\t\t\t\tGet informations about an email and check if has been compromised")

    print("\n--dns command:")
    print("\t --domain FQDN\t\tScan and get domain information")
    print("\t --host HOST\t\tScan and get host information")
    print("\t --ip IP\t\tScan and get IP information")

    print("\n--hashfile command:")
    print("\t --md5 FILE\t\tGet the MD5 hash of the file")
    print("\t --sha1 FILE\t\tGet the SHA256 of the file")
    print("\t --sha256 FILE\t\tGet the SHA256 of the file")
    print("\t --sha384 FILE\t\tGet the SHA384 of the file")
    print("\t --sha512 FILE\t\tGet the SHA512 of the file")

    print("\n--email command")
    print("\t --emailrep\t\tGet the email reputation report")

    print("\n--mac command")
    print("--mac MAC\t\t\tGet mac information")
    print("--macdb\t\t\t\tUpdate the OUI database")

def readConfigFile(config):
    """
        This function read the config file
    """
    data = {}
    try:
        with open(config, 'r') as f:
            lines = f.readlines()

        # Split each line into te dictionary
        for line in lines:
            l = line.split(":")
            lineParsed = l[1].replace(" ", "")
            lineParsed = lineParsed.replace("\n", "")
            data[l[0]] = lineParsed

    except FileNotFoundError:
        return None
    return data

def main():
    args = checkArguments()

    if not args.config:
        usage()
        exit(1);

    # Read the config file
    config = readConfigFile(args.config)
    if config is None:
        print("Failed to read the config file")
        exit(0)

    report = dict()

    # Analyse DNS
    if args.dns:
        if args.domain:
            _parsingDomain(config, args.domain, report)
        if args.host:
            _parsingHost(config, args.host, report)
        if args.ip:
            _parsingIP(config, args.ip, report)
    # Analyse hash file
    if args.hashfile:
        h = Hash()
        dispatcher = {
           "MD5": h.hashMd5,
           "SHA1": h.hashSha1,
           "SHA256": h.hashSha256,
           "SHA384": h.hashSha384,
           "SHA512": h.hashSha512,
        }
        if args.md5:
            hashType = "MD5"
            filename = args.md5
        if args.sha1:
            hashType = "SHA1"
            filename = args.sha1
        if args.sha256:
            hashType = "SHA256"
            filename = args.sha256
        if args.sha384:
            hashType = "SHA384"
            filename = args.sha384
        if args.sha512:
            hashType = "SHA512"
            filename = args.sha512
        
        if not exists(filename):
            print(f"File {filename} do not exist")
        else:
            res = dispatcher[hashType](filename)
            print(f"{hashType} hash: {res}")

            if args.scanvt:
                _parsingHash(config, res, report)

    # Analyse the hash
    if args.hash:
        _parsingHash(config, args.hash, report)

    # Analyse the email
    if args.email:
        if args.emailrep:
            _parsingEmail(config, args.emailrep)

    # Analyse mac address
    if args.macdb:
        _parseMACAddress(mac=None, db=True)
    if args.mac:
        _parseMACAddress(mac=args.mac)

def _parseMACAddress(mac=None, db=False):
    macchecker = MACChecker()
    report = dict()
    if db: 
        report['db'] = macchecker.updateOUIDb()

    if mac is not None:
        report['mac'] = macchecker.parseMACAddress(mac)

    print("----------------------------")
    print("|        MAC report        |")
    print("----------------------------")

    if 'db' in report:
        print(f"The update of the OUI db: {report['db']['success']}") 

    if 'mac' in report:
        print(f"MAC address: {report['mac']['hw']}")
        print(f"OUI: {report['mac']['oui']}")
        print(f"Vendor: {report['mac']['vendor']}")

def _parsingEmail(config, email):
    # Check if the email specified is correct
    regex = r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,7}\b'
    if not search(regex, email):
        print("Please, specify a valid email address")
        return

    report = dict()
    emailChecker = EmailChecker(config['api_key_emailrep'], email)
    report['emailrep'] = emailChecker.reportEmailRep()

    print("----------------------------")
    print("|     Email reputation     |")
    print("----------------------------")

    if 'error' in report['emailrep']:
        print(f"Error: {report['emailrep']['error']}")
        return

    emailrep = report['emailrep']
    print(f"Reputation: {emailrep['reputation']}")
    print(f"Suspicious: {emailrep['suspicious']}")

def _parsingHash(config, h, report):
    report = dict() 
    vt = VT(config['api_key_vt'])

    vt.getRateFromHash(h, report)
    print("\n----------------------")
    print("|     VirusTotal     |")
    print("----------------------")

    if 'error' in report:
        print("Error, can not find any informations")
        return

    # File
    print("\nFile information:")
    f = report['results']['file']
    print(f"\tMD5: {f['md5']}")
    print(f"\tSHA1: {f['sha1']}")
    print(f"\tMagic: {f['magic']}")
    print(f"\tFile type: {f['filetype']}")
    print(f"\tSize: {f['size']}")
    print(f"\tExtension: {f['extension']}")
    print(f"\tFirst seen: {f['first_seen']}")
    print(f"\tLast analysis: {f['last_analysis']}")

    # Vendors
    vendors = report['results']['vendors']
    print("\nVendors:")
    print(f"\tTotal vendors: {vendors['total_vendors']}")
    print(f"\tCategorized as malicious: {vendors['total_malicious']}\n")
    for key in vendors['result']:
        print(f"\t{key}: {vendors['result'][key]}")

def _parsingIP(config, ip, report):
    # Check if it's an IP address
    vt = VT(config['api_key_vt'])
    try:
        ip_obj = ipaddress.ip_address(ip)
    except ValueError:
        print("The IP is not valid")
        return

    report['ip'] = dict()
    report['ip'] = vt.getIPVirusTotal(ip)

    if 'error' not in report['ip']:
        print("----------------------------")
        print("|      IP Informations     |")
        print("----------------------------")
        for vt in VT_ATTRIBUTES_MAPPING.keys():
            try:
                vtAttributes = VT_ATTRIBUTES_MAPPING[vt]
                if 'date' in vtAttributes:
                    value = datetime.fromtimestamp(int(report['ip'][vt]))
                else:
                    value = report['ip'][vt]
                print(f"{vt}: {value}")
            except KeyError:
                print(f"Cannot find the key {vt}")

def _parsingHost(config, fqdn, report):
    vt = VT(config['api_key_vt'])
    dns = DNS(config['api_key_vt'], fqdn, {'A': 'address'})

    # Resolv and print results
    report['resolving'] = dns.resolver()
    _printDNSResolving(report['resolving'])
    
    # Print VirusTotal
    report['vt'] = dict()
    vt.getDomainReport(fqdn, report['vt'])
    _printDNSVirusTotal(report['vt'])

def _parsingDomain(config, fqdn, report):
    vt = VT(config['api_key_vt'])
    dns = DNS(config['api_key_vt'], fqdn)

    # Check if domain exist
    if not dns.checkDomainExist():
        print(f"The domain {fqdn} do not exist")
        return

    # Resolving domain
    report['resolving'] = dns.resolver()
    _printDNSResolving(report['resolving'])

    # Whois request and print the result
    report['whois'] = dns.whois()
    _printDNSWhois(report['whois'])

    # Print VirusTotal
    report['vt'] = dict()
    vt.getDomainReport(fqdn, report['vt'])
    _printDNSVirusTotal(report['vt'])

def _printDNSResolving(report):
    print("----------------------")
    print("|     resolving      |")
    print("----------------------")
    for key in report.keys():
        print(f"{key}: ")
        for entry in report[key]:
            for subkey in entry.keys():
                value = entry[subkey]
                if isinstance(value, bytes):
                    value = value.decode()
                print(f"\t{subkey}: {value}")

def _printDNSVirusTotal(report):
    print("\n----------------------")
    print("|     VirusTotal     |")
    print("----------------------")
    for key in report:
        print(f"{key}: {report[key]}")

def _printDNSWhois(report):
    print("\n----------------------")
    print("|        whois       |")
    print("----------------------")
    for key in report.keys():
        if isinstance(report[key], list):
            print(f"{key}:")
            for value in report[key]:
                print(f"\t{value}")
        else:
            print(f"{key}: {report[key]}")

if __name__ == "__main__":
    main()