"""
Sample client for the v2 API of Pentest-Tools.com.
This client starts a Web Server Scan, queries the output and writes the report in a HTML and a PDF file.
A valid API key is necessary for this program to work.

This client contains sample requests for most API methods

API Reference: https://pentest-tools.com/docs/api/v2
"""

import json
import sys
import time
import traceback
import urllib
import os
import pathlib

import requests
from typing import Literal, Union
from typing_extensions import TypedDict

USER_AGENT = os.environ.get("PTT_USER_AGENT", "ptt-cli")
API_URL = os.environ.get("PTT_API_URL", "https://app.pentest-tools.com/api/v2/")
HEADERS = {
    "Content-Type": "application/json",
    "User-Agent": USER_AGENT,
    "Accept": "*/*",
}
API_KEY = None
KEY_READ = False


def generate_api_key(xdg_config_home: str) -> dict[str, str]:
    """Generate a new freemium api key"""
    pathlib.Path(os.path.join(xdg_config_home)).mkdir(parents=True, exist_ok=True)
    config = os.path.join(xdg_config_home, "config.json")

    url = "https://pentest-tools.com/api/auth/token"
    headers = {
        "User-Agent": USER_AGENT,
        "Accept": "*/*",
    }

    response = requests.post(url, headers=headers, cookies={"credentials": "include"})
    key = response.cookies["ptt.auth.token"]

    # Also store the api key in the config file for later use
    with open(config, "w") as file:
        json.dump({"key": key}, file)

    return {"key": key}


# TODO: this doesn't work because HEADER is used as an immutable object for the functions default parameters. Changing this variable won't affect the default parameters of the function.
def init(regenerate: bool = False) -> None:
    """Init function"""
    global API_KEY
    global KEY_READ

    if regenerate or not API_KEY:
        _home = os.path.expanduser("~")

        xdg_config_home = os.environ.get("XDG_CONFIG_HOME") or os.path.join(_home, ".config")

        app_dir = os.path.join(xdg_config_home, "pentesttools")
        if windows_path := os.environ.get("LOCALAPPDATA"):
            app_dir = os.path.join(windows_path, "PentestTools", "Local")

        config = os.path.join(app_dir, "config.json")

        if not regenerate:
            # Try to read the api key from the config file
            try:
                with open(config, "r") as file:
                    config = json.load(file)
                    KEY_READ = True
            # We didn't find the key, generate the config file
            except FileNotFoundError:
                config = generate_api_key(app_dir)
        else:
            config = generate_api_key(app_dir)

        API_KEY = config["key"]

    HEADERS["Authorization"] = f"Bearer {API_KEY}"


class Tool:
    """Map the tool_id that the API knows to the tool name"""

    SUBDOMAIN_FINDER = 20
    WHOIS_LOOKUP = 30
    PORT_SCANNER = 70
    URL_FUZZER = 90
    VIRTUAL_HOSTS_FINDER = 160
    WEBSITE_SCANNER = 170
    ICMP_PING = 240
    SHARE_POINT_SCANNER = 260
    WORDPRESS_SCANNER = 270
    DRUPAL_SCANNER = 280
    JOOMLA_SCANNER = 290
    WEBSITE_RECON = 310
    NETWORK_SCANNER = 350
    SQLI_EXPLOITER = 380
    DOMAIN_FINDER = 390
    PASSWORD_AUDITOR = 400
    SSL_SCANNER = 450
    SNIPER = 490
    WAF_DETECTOR = 500
    API_SCANNER = 510
    CLOUD_SCANNER = 520
    PEOPLE_HUNTER = 530
    KUBERNETES_SCANNER = 540


# Scans

# Two common ways to start a scan is by using either `target` or `target_id`.

# `target` needs to be a simple URL, like "https://example.org". `target_id` is an
# integer you can get from the `get_targets` method.


# For both, you need:
# - tool_id: ID of the tool you want to use
# - tool_params: Options for the tool
# - target or target_id: The target you want to scan, depending on the chosen method
def start_scan(target, tool_id, tool_params, scan_params=None, api_url=API_URL, headers=HEADERS):
    """Start a scan using the given target name"""
    if scan_params is None:
        scan_params = {}

    data = {
        "tool_id": tool_id,
        "target_name": target,
        "tool_params": tool_params,
        **scan_params,
    }
    return requests.post(api_url + "scans/", headers=headers, json=data)


def start_scan_by_targetid(target_id, tool_id, tool_params, api_url=API_URL, headers=HEADERS):
    """Start a scan using the given target_id"""
    data = {"tool_id": tool_id, "target_id": target_id, "tool_params": tool_params}
    return requests.post(api_url + "scans/", headers=headers, json=data)


###Interacting with scans

# After you started a scan, through either method, you may want to interact with it. These are the most commonly
# used methods for interacting with scans after they have been started.

# A scan is identified by a `scan_id`, which can be obtained from running `GET $API/scans`,
# or `get_scans` from this client.

# You can check the status of a scan using the `get_scan_status` function.

# You can get the JSON output of a scan by calling `get_output` with a suitable `scan_id`.
# The previous feature of getting output in a chosen format has moved to the `start_scan` function, through an URL callback. TODO: example


# You may want to stop a running scan, which you can do with `stop_scan`.
# Should you want to delete a scan entirely, tou can use the `delete_scan` function.
def get_scans(
    workspace_id=None,
    target_id=None,
    status=None,
    limit=None,
    page=None,
    start_time=None,
    api_url=API_URL,
    headers=HEADERS,
):
    """Get a list of scans


            res = ptt.api.get_scans(
            workspace_id=workspace_id,
            target_id=target_id,
            status=status,
            limit=limit,
            page=page,
            start_time=start_time,
            headers=headers,
        )


    Specific parameters:
    - workspace_id  -- when set, only the scans from this workspace will be returned
                        (you can get a list of workspaces by using the `get_workspaces` operation)
    - target_id     -- when set, only the scans run on this target will be returned
                        (use `get_targets` for the target list)
    """
    data = {}
    if workspace_id is not None:
        data["workspace_id"] = workspace_id
    if target_id is not None:
        data["target_id"] = target_id
    if status is not None:
        data["status"] = status
    if limit is not None:
        data["limit"] = limit
    if page is not None:
        data["page"] = page
    if start_time is not None:
        data["start_time"] = start_time
    params = "?" + urllib.parse.urlencode(data)
    return requests.get(api_url + f"scans{params}", headers=headers)


def get_scan_status(scan_id, api_url=API_URL, headers=HEADERS):
    """Get the status of a scan"""
    return requests.get(api_url + f"scans/{scan_id}", headers=headers)


def get_output(scan_id, api_url=API_URL, headers=HEADERS):
    """Get the output of a scan"""
    return requests.get(api_url + f"scans/{scan_id}/output", headers=headers)


def get_raw_output(scan_id, api_url=API_URL, headers=HEADERS):
    """Get the raw output of a scan"""
    return requests.get(api_url + f"scans/{scan_id}/raw", headers=headers)


def stop_scan(scan_id, api_url=API_URL, headers=HEADERS):
    """Stop a running scan"""
    return requests.post(api_url + f"scans/{scan_id}/stop", headers=headers)


def delete_scan(scan_id, api_url=API_URL, headers=HEADERS):
    """Delete a scan"""
    return requests.delete(api_url + f"/scans/{scan_id}", headers=headers)


### Targets

# Although you can interact with targets manually, by inputting the URL everytime, Pentest-Tools offers facilities of working with targets.

# The simples workflow involves three functions: Add a target with `add_target`, get all targets with `get_targets` and get a single target with
# `get_target_by_id`.


# Deleting and updating targets remain, for now, an operation you can only do through the site.
def add_target(name, description="", workspace_id=None, api_url=API_URL, headers=HEADERS):
    """Add a new target

    Specific parameters:
    - name          -- the name of the target (must be a hostname, IP address or URL)
    - description   -- a short description of the target (optional)
    - workspace_id  -- the specific workspace in which to add this target (optional)
    """
    data = {"name": name}
    if len(description) > 0:
        data["description"] = description
    if workspace_id is not None:
        data["workspace_id"] = workspace_id

    return requests.post(api_url + "targets", headers=headers, json=data)


def get_targets(
    workspace_id: int = None,
    target_type: str = None,
    limit: int = None,
    page: int = None,
    api_url=API_URL,
    headers=HEADERS,
):
    """
    Get a list of targets.

    Parameters:
    - workspace_id (int, optional): ID of the workspace to show targets from.
    - type (str, optional): The type of target. Possible values: "hostname", "ip_address", "url".
    - limit (int, optional): The maximum number of targets to return.
    - page (int, optional): The page number to return.

    Returns:
    - Response object from requests.get
    """
    params: dict[str, Union[str, int]] = {}
    if workspace_id is not None:
        params["workspace_id"] = workspace_id
    if target_type is not None:
        params["type"] = target_type
    if limit is not None:
        params["limit"] = limit
    if page is not None:
        params["page"] = page

    return requests.get(api_url + "targets", headers=headers, params=params)


def get_target_by_id(target_id, api_url=API_URL, headers=HEADERS):
    """Update the description of a target

    Specific parameters:
    - target_id     -- the ID of the updated target
    """
    return requests.get(api_url + f"targets/{target_id}", headers=headers)


def delete_target(target_id, api_url=API_URL, headers=HEADERS):
    """Delete a target and its scans"""
    return requests.delete(api_url + f"targets/{target_id}", headers=headers)


### Findings
def get_findings(
    workspace_id: int = None,
    target_id: int = None,
    task_id: int = None,
    group_duplicates: bool = False,
    api_url=API_URL,
    headers=HEADERS,
):
    """Get all the findings.

    Parameters:
    - workspace_id (int, optional): ID of the workspace to show findings from.
    - target_id (int, optional): ID of the target to show findings from.
    - task_id (int, optional): ID of the scan to show findings from.
    - group_duplicates (bool, optional): Group duplicates to show only the latest finding from each group. Default is False.

    Returns:
    - Response object from requests.get
    """
    params = {}
    if workspace_id is not None:
        params["workspace_id"] = workspace_id
    if target_id is not None:
        params["target_id"] = target_id
    if task_id is not None:
        params["task_id"] = task_id
    if group_duplicates:
        params["group_duplicates"] = "true"
    else:
        params["group_duplicates"] = "false"

    return requests.get(api_url + "findings/", headers=headers, params=params)

def get_finding_information_by_id(finding_id, api_url=API_URL, headers=HEADERS):
    """Get finding information by id"""
    return requests.get(api_url + f"findings/{finding_id}", headers=headers)


### Finding templates
def get_finding_templates(
    vuln_id_type=None,
    has_ptt_exploit=None,
    limit=None,
    offset=None,
    api_url=API_URL,
    headers=HEADERS,
):
    """Get finding templates

    Specific parameters:
    - vuln_id_type      -- type of vulnerability ID (optional)
    - has_ptt_exploit   -- filter by exploit availability (optional, True/False)
    - limit             -- max number of results (optional)
    - offset            -- offset for pagination (optional)
    """
    params = {}
    if vuln_id_type is not None:
        params["vuln_id_type"] = vuln_id_type
    if has_ptt_exploit is not None:
        params["has_ptt_exploit"] = str(has_ptt_exploit)
    if limit is not None:
        params["limit"] = str(limit)
    if offset is not None:
        params["offset"] = str(offset)

    return requests.get(api_url + "public/finding_templates", headers=headers, params=params)


### Reports
def get_all_reports(api_url=API_URL, headers=HEADERS):
    """Get all reports"""
    return requests.get(api_url + "reports", headers=headers)


def create_report(
    format: Literal["pdf", "html", "json", "csv", "xlsx", "docx"],
    group_by: Literal["target", "vulnerability"],
    source: Literal["scans", "findings"],
    resources: list[int],
    filters: dict = None,
    webhook_url: str = None,
    api_url=API_URL,
    headers=HEADERS,
):
    """
    Create a report.

    Parameters:
    - format: Report format ('pdf', 'html', 'json', 'csv', 'xlsx', 'docx')
    - group_by: Grouping ('target', 'vulnerability')
    - source: Source ('scans', 'findings')
    - resources: List of scan or finding IDs
    - filters: Optional dict of filter options
    - webhook_url: Optional webhook URL for report delivery

    Returns:
    - Response object from requests.post
    """
    data = {
        "format": format,
        "group_by": group_by,
        "source": source,
        "resources": resources,
    }
    if filters is not None:
        data["filters"] = filters
    if webhook_url is not None:
        data["webhook_url"] = webhook_url

    return requests.post(api_url + "reports", headers=headers, json=data)


def get_report_information_by_id(report_id, api_url=API_URL, headers=HEADERS):
    """Get report information by id"""
    return requests.get(f"{api_url}reports/{report_id}", headers=headers)


def delete_report(report_id, api_url=API_URL, headers=HEADERS):
    """Delete a report by its ID"""
    return requests.delete(f"{api_url}reports/{report_id}", headers=headers)


def download_report(report_id, api_url=API_URL, headers=HEADERS):
    """
    Download a report by its ID.

    Parameters:
    - report_id: The ID of the report to download

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}reports/{report_id}/download", headers=headers)


### VPN profiles
def get_vpn_profiles(api_url=API_URL, headers=HEADERS):
    """
    Get all VPN profiles.

    Returns:
    - Response object from requests.get
    """
    return requests.get(api_url + "vpn_profiles", headers=headers)


### Wordlists
def get_wordlists(api_url=API_URL, headers=HEADERS):
    """
    Get all available wordlists.

    Returns:
    - Response object from requests.get
    """
    return requests.get(api_url + "wordlists", headers=headers)


def get_wordlist_information_by_id(wordlist_id, api_url=API_URL, headers=HEADERS):
    """
    Get information about a specific wordlist by its ID.

    Parameters:
    - wordlist_id: The ID of the wordlist

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}wordlists/{wordlist_id}", headers=headers)


def delete_wordlist(wordlist_id, api_url=API_URL, headers=HEADERS):
    """
    Delete a wordlist by its ID.

    Parameters:
    - wordlist_id: The ID of the wordlist to delete

    Returns:
    - Response object from requests.delete
    """
    return requests.delete(f"{api_url}wordlists/{wordlist_id}", headers=headers)


def get_wordlists_content_by_id(wordlist_id, api_url=API_URL, headers=HEADERS):
    """
    Get the contents of a specific wordlist by its ID.

    Parameters:
    - wordlist_id: The ID of the wordlist

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}wordlists/{wordlist_id}/contents", headers=headers)


### HTTP Loggers
def get_http_loggers(workspace_id=None, api_url=API_URL, headers=HEADERS):
    """
    Get all HTTP loggers.

    Parameters:
    - workspace_id: Optional workspace ID to filter loggers

    Returns:
    - Response object from requests.get
    """
    params = {}
    if workspace_id is not None:
        params["workspace_id"] = workspace_id
    return requests.get(api_url + "http_loggers", headers=headers, params=params)


def create_http_logger(label: str, workspace_id: int = None, api_url=API_URL, headers=HEADERS):
    """
    Create a new HTTP request logger.

    Parameters:
    - label: A unique label for the HTTP request logger (required)
    - workspace_id: ID of the workspace to associate the logger with (optional)

    Returns:
    - Response object from requests.post
    """
    data = {"label": label}
    if workspace_id is not None:
        data["workspace_id"] = workspace_id

    # The API expects application/x-www-form-urlencoded
    form_headers = headers.copy()
    form_headers["Content-Type"] = "application/x-www-form-urlencoded"

    return requests.post(
        api_url + "http_loggers",
        headers=form_headers,
        data=urllib.parse.urlencode(data),
    )


def get_logger_information_by_id(logger_id, api_url=API_URL, headers=HEADERS):
    """
    Get information about a specific HTTP logger by its ID.

    Parameters:
    - logger_id: The ID of the HTTP logger

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}http_loggers/{logger_id}", headers=headers)


def delete_logger(logger_id, api_url=API_URL, headers=HEADERS):
    """
    Delete an HTTP logger by its ID.

    Parameters:
    - logger_id: The ID of the HTTP logger to delete

    Returns:
    - Response object from requests.delete
    """
    return requests.delete(f"{api_url}http_loggers/{logger_id}", headers=headers)


def get_data_for_logger(logger_id, api_url=API_URL, headers=HEADERS):
    """
    Get data captured by a specific HTTP logger.

    Parameters:
    - logger_id: The ID of the HTTP logger

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}http_loggers/{logger_id}/data", headers=headers)


def clear_data_for_logger(logger_id, api_url=API_URL, headers=HEADERS):
    """
    Clear all data captured by a specific HTTP logger.

    Parameters:
    - logger_id: The ID of the HTTP logger

    Returns:
    - Response object from requests.delete
    """
    return requests.delete(f"{api_url}http_loggers/{logger_id}/data", headers=headers)


### Workspaces
def get_workspaces(api_url=API_URL, headers=HEADERS):
    """
    Get all workspaces.

    Returns:
    - Response object from requests.get
    """
    return requests.get(api_url + "workspaces", headers=headers)


def get_workspace_by_id(workspace_id, api_url=API_URL, headers=HEADERS):
    """
    Get information about a specific workspace by its ID.

    Parameters:
    - workspace_id: The ID of the workspace

    Returns:
    - Response object from requests.get
    """
    return requests.get(f"{api_url}workspaces/{workspace_id}", headers=headers)


def delete_workspace(workspace_id, api_url=API_URL, headers=HEADERS):
    """
    Delete a workspace by its ID.

    Parameters:
    - workspace_id: The ID of the workspace to delete

    Returns:
    - Response object from requests.delete
    """
    return requests.delete(f"{api_url}workspaces/{workspace_id}", headers=headers)


def create_workspace(name: str, description: str = "", api_url=API_URL, headers=HEADERS):
    """
    Create a new workspace.

    Parameters:
    - name: Name of the workspace (required)
    - description: Description of the workspace (optional)

    Returns:
    - Response object from requests.post
    """
    data = {"name": name}
    if description:
        data["description"] = description

    form_headers = headers.copy()
    form_headers["Content-Type"] = "application/x-www-form-urlencoded"

    return requests.post(
        api_url + "workspaces",
        headers=form_headers,
        data=urllib.parse.urlencode(data),
    )


def edit_workspace(
    workspace_id: int,
    name: str = None,
    description: str = None,
    api_url=API_URL,
    headers=HEADERS,
):
    """
    Edit an existing workspace.

    Parameters:
    - workspace_id: ID of the workspace to edit (required)
    - name: New name for the workspace (optional)
    - description: New description for the workspace (optional)

    Returns:
    - Response object from requests.put
    """
    data = {}
    if name is not None:
        data["name"] = name
    if description is not None:
        data["description"] = description

    return requests.put(
        f"{api_url}workspaces/{workspace_id}",
        headers=headers,
        json=data,
    )


class Finding(TypedDict):
    name: str
    source: int
    target_name: str
    test_description: str
    vuln_description: str
    vuln_risk_level: int
    vuln_recommendation: str
    cvss_v3: float
    vuln_evidence: str


def import_findings(
    findings: list[Finding],
    workspace_id: int = None,
    api_url=API_URL,
    headers=HEADERS,
):
    """
    The import_findings API call.

    Parameters:
    - findings: list of FindingDict
    - workspace_id: int (optional)
    """
    data: dict[str, Union[list, int]] = {}
    if workspace_id is not None:
        data["workspace_id"] = workspace_id

    data["findings"] = findings

    return requests.post(
        api_url + "findings/import",
        headers=headers,
        json=data,
    )


if __name__ == "__main__":
    # `tool_params` is specific to the tool
    # Here we do a light scan with the Web Server Scanner
    tool_id = Tool.WEBSITE_SCANNER
    tool_params = {"scan_type": "light"}
    target = "http://demo.pentest-tools.com/webapp/"

    # Start the scan
    res = start_scan(target, tool_id, tool_params)
    try:
        res_json = res.json()
    except requests.exceptions.JSONDecodeError:
        print(traceback.format_exc())
        print(res.text)
        sys.exit(1)

    # Get the new `scan_id`
    if "data" in res_json and "created_id" in res_json["data"]:
        scan_id = res_json["data"]["created_id"]
        print("Started scan %i" % scan_id)
    else:
        print("Scan could not start")
        print(f"Status: {res_json['status']}, message: {res_json['message']}")
        sys.exit(1)

    # Poll periodically to check if the scan is finished
    while True:
        time.sleep(2)

        # Get the status of our scan
        status = get_scan_status(scan_id)
        status_name = status.json()["data"]["status_name"]

        if status_name == "finished":
            print("Scan status: %s" % res_json["data"])
            # Get the HTML report and write it to a file
            print("Getting JSON report")
            res = get_output(scan_id)
            output_json = res.json()

            with open("report.json", "w") as file:
                json.dump(output_json, file)

            print("JSON report written to file")
            break
