# Author: perp
# Description: Sort nuclei json into folders
# Usage: nukesort scan.json results

import json
import os
import shutil
import sys

# Check arguments
if len(sys.argv) < 3:
    print("You are missing an argument! (Requires: json & folder)")
    exit()

# Json & folder
scan_json = str(sys.argv[1])
scan_folder = str(sys.argv[2])

# Path doesn't exist, make
if not os.path.isdir(scan_folder):
    os.makedirs(f"{scan_folder}/domains")

# Path exists, remove, remake
elif os.path.isdir(scan_folder):
    shutil.rmtree(f"{scan_folder}/domains", ignore_errors=True)
    os.makedirs(f"{scan_folder}/domains")

# Open json file
with open(scan_json, "r") as f:
    # Load JSON
    lines = json.loads(f.read())

    # Go through each line
    for line in range(len(lines)):
        # Get domain & vulnerability
        domain = lines[line]["host"].replace("http://", "").replace("https://", "").replace("www", "")
        vulnerability = lines[line]["template-id"]

        # Create domain folder
        if not os.path.exists(f"{scan_folder}/domains/{domain}"):
            os.mkdir(f"{scan_folder}/domains/{domain}")

        # Create vulnerability folder
        if not os.path.exists(f"{scan_folder}/domains/{domain}/{vulnerability}"):
            os.mkdir(f"{scan_folder}/domains/{domain}/{vulnerability}")

        # Result found
        if "extracted-results" in lines[line]:
            # Open vulnerability text file
            with open(f"{scan_folder}/domains/{domain}/{vulnerability}/{vulnerability}.txt", "a") as vuln:
                # Write vulnerability
                results = lines[line]["extracted-results"]
                for result in results:
                    vuln.write(result+"\n")

        # Open vulnerability json file
        with open(f"{scan_folder}/domains/{domain}/{vulnerability}/{vulnerability}.json", "a") as vulns:
            # Write json
            json.dump(lines[line], vulns, indent=4)