#!/usr/bin/env python
import argparse
import json
import uuid
from termcolor import colored

"""
Script to generate a CSV and a MINDMAP (via mermaid) representations
from the output file of DNSX that gathered ALL DNS records types
from a collection of domain names.

DNSX executed command:

    dnsx -silent -duc -l list_of_domain_names.txt -recon -r 8.8.8.8,1.1.1.1 -o data.json -json

Sources:
    https://github.com/projectdiscovery/dnsx
    https://mermaid.js.org/syntax/mindmap.html
    https://www.cloudflare.com/learning/dns/dns-records/

Dependencies:
    pip install colorama termcolor
"""


def load_records(dnsx_file):
    dns_record_types = ["cname"]
    # KEY is the IP
    # VALUE is a dict for which: KEY is the DNS record type (except the 'A' one) and VALUE is an array of the values for the records
    records = {}
    # DNSX JSON file is not well formed because each line is a complete JSON object.
    with open(dnsx_file, mode="r", encoding="utf-8") as f:
        lines = f.read().splitlines()
    for line in lines:
        json_data = json.loads(line)
        domain = json_data["host"]
        for ip in json_data["a"]:
            if ip not in records:
                records[ip] = {}
            if "domain" not in records[ip]:
                records[ip]["domain"] = []
            if domain not in records[ip]["domain"]:
                records[ip]["domain"].append(domain)
            for dns_record_type in dns_record_types:
                if dns_record_type in json_data:
                    if dns_record_type not in records[ip]:
                        records[ip][dns_record_type] = []
                    records[ip][dns_record_type].extend(json_data[dns_record_type])
                    records[ip][dns_record_type] = list(set(records[ip][dns_record_type]))
    return records


def generate_csv_representation(records, output_base_filename):
    entries = []
    for domain, dns_record_types in records.items():
        for dns_record_type in dns_record_types:
            for value in dns_record_types[dns_record_type]:
                entries.append(f"{domain},{dns_record_type.upper()},{value}")
    entries.sort()
    with open(f"{output_base_filename}.csv", mode="w", encoding="utf-8") as f:
        f.write("IP,DNS-Record-Type,Value\n")
        f.write("\n".join(entries))


def generate_mindmap_representation(records, output_base_filename):
    separator = "  "
    mermaid_content = "mindmap\n"
    mermaid_content += separator + "root{{IP}}\n"
    for ip, dns_record_types in records.items():
        mermaid_content += f'{separator*2}[{ip}]\n'
        for dns_record_type in dns_record_types:
            id = uuid.uuid4()
            mermaid_content += f'{separator*3}{id}("{dns_record_type.upper()} ({len(dns_record_types[dns_record_type])})")\n'
            for value in dns_record_types[dns_record_type]:
                mermaid_content += f'{separator*4}{value}\n'
    with open(f"{output_base_filename}.md", mode="w", encoding="utf-8") as f:
        f.write(mermaid_content)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Script to generate a CSV and a MINDMAP (via mermaid) representations from the output file of DNSX that gathered all DNS records types from a collection of domain names.")
    parser.add_argument("-f", action="store", dest="dnsx_file", help="JSON file coming from the DNSX executed commands", required=True)
    parser.add_argument("-o", action="store", dest="output_base_filename", help="Base name of the output file", required=False, default="/tmp/representations")
    args = parser.parse_args()
    output_base_filename = args.output_base_filename
    records = load_records(args.dnsx_file)
    print(colored(f"[+] Generate the CSV and Markdown representations to '{output_base_filename}.(csv|md)' files ...", "yellow"))
    generate_csv_representation(records, output_base_filename)
    generate_mindmap_representation(records, output_base_filename)
    print(colored(f"[V] Done (Site https://mermaid.live can be used to view the mindmap).", "green"))
