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

"""
Script to generate a CSV and a MINDMAP (via mermaid) representations
from the output file of DNSX that gathered the PTR DNS records
from a collection of IP addresses.

DNSX executed command:

    dnsx -silent -duc -l list_of_ip_addresses_or_cidr.txt -resp-only -ptr -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/dns-ptr-record/
    https://aws.amazon.com/what-is/cidr

Dependencies:
    pip install colorama termcolor
"""


def load_records(dnsx_file):
    # KEY is the IP and the VALUE is the list of domains the IP is refering to.
    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)
        ip = json_data["host"]
        domains = json_data["ptr"]
        records[ip] = domains
    return records


def generate_csv_representation(records, output_base_filename):
    entries = []
    for ip, domains in records.items():
        for domain in domains:
            entries.append(f"{ip},{domain}")
    entries.sort()
    with open(f"{output_base_filename}.csv", mode="w", encoding="utf-8") as f:
        f.write("IP,Domain\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, domains in records.items():
        mermaid_content += f'{separator*2}{ip}("{ip} ({len(domains)})")\n'
        for domain in domains:
            mermaid_content += f"{separator*3}{domain}\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 the PTR DNS records from a collection of IP addresses.")
    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"))
