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

from argparse import ArgumentParser
from config import VT_ATTRIBUTES_MAPPING
from vt import VT
from dnsinformations import DNSInformations as DNS
import ipaddress
from datetime import datetime
from hashing import Hash
from os.path import exists


def checkArguments():
    parser = ArgumentParser(description="baoSOC")
    parser.add_argument('-c', '--config', help='Config file')
    parser.add_argument('--dns', help='Get domain name information', action="store_true")
    # For dns command
    parser.add_argument('--dnsattacks', help='Parse DNS pcap file')
    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('--hash', 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')

    return parser.parse_args()

def usage():
    print("------------------------------")
    print("|           baoSOC           |")
    print("------------------------------\n")
    print("A tool for SOC analyst\n")
    print("Usage: main.py [COMMAND]")
    print("-c PATH, --config PATH\t\tConfig file - mandatory")
    print("--hash FILE\t\t\tHash the file and check in VirusTotal")
    print("--dns \t\t\t\tGet information regarding the domain with whois and VirusTotal")
    print("--dnsattacks FILE\t\tParse the DNS pcap file and identify some DNS attacks")

    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--hash 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")

def mainMenu():
    print("\n           baoSOC           ")
    print(" What would you like to do? ")
    print("\n OPTION 1: Sanitise URL For emails ")
    print(" OPTION 2: Decoders (PP, URL, SafeLinks) ")
    print(" OPTION 3: Reputation Checker")
    print(" OPTION 4: DNS Tools")
    print(" OPTION 5: Hashing Function")
    print(" OPTION 6: Phishing Analysis")
    print(" OPTION 7: URL scan")
    print(" OPTION 9: Extras")
    print(" OPTION 0: Exit Tool")

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:
            pass
        if args.ip:
            _parsingIP(config, args.ip, report)
    # Analyse hash
    if args.hash:
        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}")

        res = "6548eec09f4d8bc6514bee3e5452541c"
        if args.scanvt:
            _parsingHash(config, res, report)

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

    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'])
    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 _parsingDomain(config, fqdn, report):
    vt = VT(config['api_key'])
    dns = DNS(config['api_key'], fqdn)

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

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

    print("\n----------------------")
    print("|     VirusTotal     |")
    print("----------------------")
    report['vt'] = dict()
    vt.getDomainReport(fqdn, report['vt'])
    report_vt = report['vt']
    for key in report_vt:
        print(f"{key}: {report_vt[key]}")


if __name__ == "__main__":
    main()