#!/usr/bin/python

"""
Sample API client for 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/api-reference
"""

import json
import time
import traceback
import base64
import requests

def start_scan(api_url, tool_id, target, tool_params):
    """ Start a scan using the given target name

    Specific parameters:
    - tool_id       -- the ID of the tool to start (check the API documentation for the possible values)
    - target        -- the target in string format (e.g. "example.org", "https://example.org")
    - tool_params   -- the parameters of the tool (check the API documentation for the specific formats)
    """
    data = {
        "op":           "start_scan",
        "tool_id":      tool_id,
        "target":       target,
        "tool_params":  tool_params
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def start_scan_by_targetid(api_url, tool_id, target_id, tool_params):
    """ Start a scan using the given target_id

    Specific parameters:
    - tool_id       -- the ID of the tool to start (check the API documentation for the possible values)
    - target_id     -- the target in integer format (use `get_targets` for the target list)
    - tool_params   -- the parameters of the tool (check the API documentation for the specific formats)
    """
    data = {
        "op": "start_scan_by_targetid",
        "tool_id": tool_id,
        "target_id": target_id,
        "tool_params": tool_params
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def get_output(api_url, scan_id, output_format):
    """ Get the output of a scan

    Specific parameters:
    - scan_id       -- the scan_id of the requested scan (use `get_scans` for the scan list)
    - output_format -- the format in which to get the output of the scan ("html"/"json"/"pdf")
    """
    data = {
        "op": "get_output",
        "scan_id": scan_id,
        "output_format": output_format
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def stop_scan(api_url, scan_id):
    """ Stop a running scan

    Specific parameters:
    - scan_id -- the scan_id of the requested scan (use `get_scans` for the scan list)
    """
    data = {
        "op": "stop_scan",
        "scan_id": scan_id
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def get_scan_status(api_url, scan_id):
    """ Get the status of a scan

    Specific parameters:
    - scan_id -- the scan_id of the requested scan (use `get_scans` for the scan list)
    """
    data = {
        "op": "get_scan_status",
        "scan_id": scan_id
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def get_scans(api_url, limit, workspace_id=None, target_id=None):
    """ Get a list of scans

    Specific parameters:
    - limit         -- maximum number of returned scans
    - 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 = {
        "op": "get_scans",
        "limit": limit
    }
    if workspace_id is not None:
        data["workspace_id"] = workspace_id
    if target_id is not None:
        data["target_id"] = target_id
    res = requests.post(api_url, data=json.dumps(data))
    return res


def get_targets(api_url, limit=0):
    """ Get a list of targets

    Specific parameters:
    - limit -- maximum number of returned targets (up to 1000 will be returned if it is set to 0)
    """
    data = {
        "op": "get_targets",
        "limit": limit
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


def add_target(api_url, name, description="", workspace_id=None):
    """ 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 = {
        "op": "add_target",
        "name": name
    }
    if len(description) > 0:
        data["description"] = description
    if workspace_id is not None:
        data["workspace_id"] = workspace_id

    res = requests.post(api_url, data=json.dumps(data))
    return res


def update_target_description(api_url, target_id, description):
    """ Update the description of a target

    Specific parameters:
    - target_id     -- the ID of the updated target
    - description   -- the new description of the target
    """
    data = {
        "op": "update_target_description",
        "target_id": target_id,
        "description": description
    }
    res = requests.post(api_url, data=json.dumps(data))
    return res


if __name__ == '__main__':

    key = "xxxxxxxxxxxxxx"                              #    <-- Put your API key here
    api_url = "https://pentest-tools.com/api?key=" + key

    # These values are specific to each tool
    # In this case we want to use the Web Server Scanner
    # See the API Reference for more details
    tool_id = 170
    tool_params = {
        "scan_type": "quick"
    }
    target = "http://demo.pentest-tools.com/webapp/"

    # Start the scan
    res = start_scan(api_url, tool_id, target, tool_params)
    try:
        res_json = json.loads(res.text)
    except:
        print(traceback.format_exc())
        print(res.text)


    # Check the status of the operation
    if res_json["op_status"] == "success":

        # This is the id of the new scan
        scan_id = res_json["scan_id"]

        print("Started scan %i" % scan_id)


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

            # Get the status of our scan
            res = get_scan_status(api_url, scan_id)
            res_json = json.loads(res.text)

            if res_json["op_status"] == "success":

                print("Scan status: %s" % res_json["scan_status"])

                if res_json["scan_status"] != "waiting" and res_json["scan_status"] != "running":

                    # Get the HTML report and write it to a file
                    print("Getting HTML report")
                    res = get_output(api_url, scan_id, "html")
                    res_json = json.loads(res.text)
                    output_html = res_json["scan_output"]["output_html"]

                    with open("report.html", "w") as f:
                        f.write(output_html)
                    print("HTML report written to file")


                    # Get the PDF report and write it to a file
                    print("Getting PDF report")
                    res = get_output(api_url, scan_id, "pdf")
                    res_json = json.loads(res.text)
                    output_pdf = res_json["scan_output"]["output_pdf"]

                    with open("report.pdf", "wb") as f:
                        f.write(base64.b64decode(output_pdf))
                    print("PDF report written to file")

                    break

            else:
                print("Operation get_scan_status failed because: %s. %s" % (res_json["error"], res_json["details"]))
                break

    else:
        print("Operation start_scan failed because: %s. %s" % (res_json["error"], res_json["details"]))