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

def tunnelingDNSAttacks(report, queries, dnsbl):
    """
        This function identify if we have a DNS tunneling
        We can identify the payload size (512 bytes) and for a specific domain, how many occurs it appear.
        For the payload size, according to the RFC 1035: https://datatracker.ietf.org/doc/html/rfc1035#section-2.3.4
        The maximum UDP message is 512 or less.
        If we detect a lot of request to a specific domain, that can be a DNS tunneling
    """
    report['blackDomain'] = _blackDomain(queries, dnsbl)
    report['oversized'] = _oversized(queries)

def _blackDomain(queries, dnsbl) -> list:
    """"
    This function identify if a dns domain is in a blacklist
    Return the report
    """
    report = dict()

    # Read the black domain list
    dataDnsbl = list()
    with open(dnsbl, 'r') as f:
        for line in f:
            dataDnsbl.append(line.replace("\n", ""))

    for query in queries:
        domain = _splitDomain(query['query'])

        # Check if in dns blacklist
        if domain in dataDnsbl:
            if domain not in report:
                report[domain] = list()
                report[domain].append(query['src'])
            else:
                if query['src'] not in report[domain]:
                    report[domain].append(query['src'])
    return report

def _oversized(queries) -> dict:
    """
        This function analyzed the payload size and if it's oversize, that can be a DNS tunneling
    """
    report = dict()

    for query in queries:
        if query['len'] > 220:
            domain = _splitDomain(query['query'])

            if domain not in report:
                report[domain] = list()
                report[domain].append({
                    'source': query['src'],
                    'len': query['len'],
                })
            else:
                if query['src'] not in report[domain]:
                    report[domain].append({
                        'source': query['src'],
                        'len': query['len'],
                    })
                

    return report

def _splitDomain(query) -> str:
    """
        This function split the query for identifying the domain
    """
    # Split to get the country code and the TLD
    domainSplitted = str(query).split(".")
    l = len(domainSplitted)
    tld = domainSplitted[l - 3:l - 2]
    cc = domainSplitted[l - 2 :l - 1]
    try:
        domain = f"{tld[0]}.{cc[0]}"
    except IndexError:
        print(f"Failed to parse the domain {query}")
        domain = None

    return domain