#!/usr/bin/env python3
"""Write recon crawl batch 2 (upload hunter + admin surface), test, commit, push."""
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]

ADMIN_SURFACE = r'''package recon

import (
	"fmt"
	"net/http"
	"sort"
	"strings"
)

var adminSurfacePaths = []string{
	"/wp-admin",
	"/wp-admin/",
	"/admin",
	"/admin/",
	"/admin/login",
	"/administrator",
	"/api",
	"/api/",
	"/api/v1",
	"/graphql",
	"/graphql/",
	"/swagger",
	"/swagger/",
	"/swagger/index.html",
	"/swagger-ui",
	"/swagger-ui/",
	"/actuator",
	"/actuator/",
	"/actuator/health",
	"/.env",
	"/.env.local",
	"/server-status",
	"/server-status/",
}

func ProbeAdminSurface(host string, port int, scheme string) []AdminSurfaceFinding {
	scheme = normalizeScheme(scheme, port)
	if port <= 0 {
		port = defaultPortForScheme(scheme)
	}
	base := fmt.Sprintf("%s://%s", scheme, joinHostPort(host, port))

	seen := map[string]bool{}
	var out []AdminSurfaceFinding
	for _, path := range adminSurfacePaths {
		key := strings.ToLower(path)
		if seen[key] {
			continue
		}
		seen[key] = true

		rawURL := strings.TrimRight(base, "/") + path
		status, _, err := fetchPage(rawURL)
		if err != nil {
			continue
		}
		signal := adminSurfaceSignal(status)
		if signal == "" {
			continue
		}
		out = append(out, AdminSurfaceFinding{
			Path:       path,
			URL:        rawURL,
			StatusCode: status,
			Signal:     signal,
		})
	}

	sort.Slice(out, func(i, j int) bool {
		if out[i].Signal != out[j].Signal {
			return out[i].Signal == "green"
		}
		return out[i].Path < out[j].Path
	})
	return out
}

func adminSurfaceSignal(status int) string {
	switch status {
	case http.StatusOK:
		return "green"
	case http.StatusUnauthorized, http.StatusForbidden:
		return "gray"
	default:
		return ""
	}
}
'''

UPLOAD_HUNTER = r'''package recon

import (
	"net/url"
	"regexp"
	"sort"
	"strings"
)

var (
	openUploadPathRe = regexp.MustCompile(`(?i)(/api/[^\s"'` + "`" + r`<>]*upload|/upload[^\s"'` + "`" + r`<>]*|/v\d+/upload)`)
	dragDropClassRe  = regexp.MustCompile(`(?i)(dropzone|drop-zone|file-drop|drag-drop|fileupload)`)
	jsUploadHintRe   = regexp.MustCompile(`(?i)(multipart/form-data|formdata\s*\(|type\s*:\s*['"]file['"]|/api/[^\s"'` + "`" + r`<>]*upload|\.upload\s*\(|dropzone)`)
)

func collectUploadFromPage(pageURL string, fileInputs, multipart []FormFinding) []UploadHunterFinding {
	var out []UploadHunterFinding
	for _, f := range fileInputs {
		if !f.HasFile {
			continue
		}
		out = append(out, UploadHunterFinding{
			PageURL: pageURL,
			Target:  resolveUploadTarget(pageURL, f.Action),
			Source:  "file_input",
			Method:  f.Method,
		})
	}
	for _, f := range multipart {
		if !f.Multipart {
			continue
		}
		out = append(out, UploadHunterFinding{
			PageURL: pageURL,
			Target:  resolveUploadTarget(pageURL, f.Action),
			Source:  "multipart",
			Method:  f.Method,
		})
	}
	return out
}

func detectDragDropZones(pageURL, body string) []UploadHunterFinding {
	root, err := htmlParseRoot(body)
	if err != nil {
		return detectDragDropFromText(pageURL, body)
	}
	var out []UploadHunterFinding
	var walk func(*htmlNode)
	walk = func(n *htmlNode) {
		if n.tag != "" {
			cls := strings.ToLower(n.attr("class"))
			id := strings.ToLower(n.attr("id"))
			dropAttr := strings.ToLower(n.attr("data-dropzone"))
			if dragDropClassRe.MatchString(cls) || dragDropClassRe.MatchString(id) || dropAttr != "" {
				target := n.attr("data-upload-url")
				if target == "" {
					target = n.attr("action")
				}
				out = append(out, UploadHunterFinding{
					PageURL: pageURL,
					Target:  resolveUploadTarget(pageURL, target),
					Source:  "drag_drop",
				})
			}
		}
		for _, c := range n.children {
			walk(c)
		}
	}
	walk(root)
	if len(out) == 0 {
		return detectDragDropFromText(pageURL, body)
	}
	return out
}

func detectDragDropFromText(pageURL, body string) []UploadHunterFinding {
	lower := strings.ToLower(body)
	if !strings.Contains(lower, "dropzone") && !strings.Contains(lower, "drag") {
		return nil
	}
	if !strings.Contains(lower, "upload") && !strings.Contains(lower, "file") {
		return nil
	}
	return []UploadHunterFinding{{
		PageURL: pageURL,
		Source:  "drag_drop",
	}}
}

func extractScriptSrc(pageURL, body string) []string {
	root, err := htmlParseRoot(body)
	if err != nil {
		return nil
	}
	var srcs []string
	var walk func(*htmlNode)
	walk = func(n *htmlNode) {
		if n.tag == "script" {
			if src := strings.TrimSpace(n.attr("src")); src != "" {
				srcs = append(srcs, src)
			}
		}
		for _, c := range n.children {
			walk(c)
		}
	}
	walk(root)
	_ = pageURL
	return srcs
}

func scanJSForUpload(jsURL, body string) []UploadHunterFinding {
	if !jsUploadHintRe.MatchString(body) {
		return nil
	}
	target := jsURL
	if m := openUploadPathRe.FindString(body); m != "" {
		target = m
	}
	return []UploadHunterFinding{{
		PageURL: jsURL,
		Target:  target,
		Source:  "js",
	}}
}

func isScriptAsset(ref string) bool {
	ref = strings.ToLower(strings.TrimSpace(ref))
	return strings.HasSuffix(ref, ".js") || strings.Contains(ref, ".js?")
}

func resolveUploadTarget(pageURL, action string) string {
	action = strings.TrimSpace(action)
	if action == "" {
		if u, err := url.Parse(pageURL); err == nil {
			return u.Path
		}
		return pageURL
	}
	if strings.HasPrefix(action, "http://") || strings.HasPrefix(action, "https://") {
		return action
	}
	if abs, err := resolveSameOrigin(pageURL, action); err == nil {
		return abs
	}
	return action
}

func rankUploadFindings(base string, in []UploadHunterFinding) []UploadHunterFinding {
	if len(in) == 0 {
		return nil
	}
	seen := map[string]UploadHunterFinding{}
	for _, f := range in {
		key := strings.ToLower(f.PageURL + "|" + f.Target + "|" + f.Source)
		if prev, ok := seen[key]; ok {
			if scoreUploadFinding(f) > scoreUploadFinding(prev) {
				seen[key] = f
			}
			continue
		}
		seen[key] = f
	}
	out := make([]UploadHunterFinding, 0, len(seen))
	for _, f := range seen {
		f = tagUploadFinding(base, f)
		f.Score = scoreUploadFinding(f)
		out = append(out, f)
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Score != out[j].Score {
			return out[i].Score > out[j].Score
		}
		return out[i].Target < out[j].Target
	})
	return out
}

func tagUploadFinding(base string, f UploadHunterFinding) UploadHunterFinding {
	target := f.Target
	if target == "" {
		target = f.PageURL
	}
	path := target
	if u, err := url.Parse(target); err == nil && u.Path != "" {
		path = u.Path
	}
	if openUploadPathRe.MatchString(path) || openUploadPathRe.MatchString(target) {
		f.Tags = appendUniqueTag(f.Tags, "open_api")
	}
	probeURL := target
	if !strings.HasPrefix(probeURL, "http://") && !strings.HasPrefix(probeURL, "https://") {
		if abs, err := resolveSameOrigin(base, probeURL); err == nil {
			probeURL = abs
		} else {
			probeURL = strings.TrimRight(base, "/") + "/" + strings.TrimLeft(probeURL, "/")
		}
	}
	status, _, err := fetchPage(probeURL)
	if err == nil {
		f.StatusCode = status
		if status == 200 {
			f.Tags = appendUniqueTag(f.Tags, "no_auth")
		}
	}
	return f
}

func scoreUploadFinding(f UploadHunterFinding) int {
	score := 10
	switch f.Source {
	case "multipart":
		score += 30
	case "file_input":
		score += 25
	case "drag_drop":
		score += 20
	case "js":
		score += 15
	}
	for _, tag := range f.Tags {
		switch tag {
		case "no_auth":
			score += 40
		case "open_api":
			score += 25
		}
	}
	if f.StatusCode == 200 {
		score += 10
	}
	return score
}

func appendUniqueTag(tags []string, tag string) []string {
	for _, t := range tags {
		if t == tag {
			return tags
		}
	}
	return append(tags, tag)
}
'''

TEST_APPEND = r'''

func TestUploadHunterMultipartAndJS(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/":
			w.Write([]byte(` + "`" + r`<html><body>
<form action="/api/upload" method="post" enctype="multipart/form-data">
<input type="file" name="payload">
</form>
<div class="dropzone" data-upload-url="/api/upload"></div>
<script src="/static/upload.js"></script>
</body></html>` + "`" + r`))
		case "/static/upload.js":
			w.Header().Set("Content-Type", "application/javascript")
			w.Write([]byte(` + "`" + r`fetch("/api/upload",{method:"POST",body:new FormData()})` + "`" + r`))
		case "/api/upload":
			w.WriteHeader(http.StatusOK)
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	u, err := url.Parse(srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	port := 80
	if p := u.Port(); p != "" {
		port = atoi(p)
	}

	SetFetchPageHook(func(rawURL string) (int, string, error) {
		resp, err := http.Get(rawURL)
		if err != nil {
			return 0, "", err
		}
		defer resp.Body.Close()
		body, _ := readBodyLimited(resp.Body, maxHTMLBytes)
		return resp.StatusCode, body, nil
	})
	t.Cleanup(func() { SetFetchPageHook(nil) })

	report, err := Crawl(u.Hostname(), port, u.Scheme, []string{"/"})
	if err != nil {
		t.Fatal(err)
	}
	if len(report.UploadHunter) == 0 {
		t.Fatal("expected upload hunter findings")
	}
	if report.UploadHunter[0].Score <= 0 {
		t.Fatalf("score=%d", report.UploadHunter[0].Score)
	}
	hasOpenAPI := false
	hasNoAuth := false
	for _, f := range report.UploadHunter {
		for _, tag := range f.Tags {
			if tag == "open_api" {
				hasOpenAPI = true
			}
			if tag == "no_auth" {
				hasNoAuth = true
			}
		}
	}
	if !hasOpenAPI {
		t.Fatalf("missing open_api tag: %+v", report.UploadHunter)
	}
	if !hasNoAuth {
		t.Fatalf("missing no_auth tag: %+v", report.UploadHunter)
	}
}

func TestProbeAdminSurfaceSignals(t *testing.T) {
	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/admin", "/admin/":
			w.WriteHeader(http.StatusForbidden)
		case "/graphql", "/graphql/":
			w.WriteHeader(http.StatusUnauthorized)
		case "/swagger", "/swagger/":
			w.WriteHeader(http.StatusOK)
			w.Write([]byte("swagger-ui"))
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	u, err := url.Parse(srv.URL)
	if err != nil {
		t.Fatal(err)
	}
	port := 80
	if p := u.Port(); p != "" {
		port = atoi(p)
	}

	SetFetchPageHook(func(rawURL string) (int, string, error) {
		resp, err := http.Get(rawURL)
		if err != nil {
			return 0, "", err
		}
		defer resp.Body.Close()
		body, _ := readBodyLimited(resp.Body, maxHTMLBytes)
		return resp.StatusCode, body, nil
	})
	t.Cleanup(func() { SetFetchPageHook(nil) })

	findings := ProbeAdminSurface(u.Hostname(), port, u.Scheme)
	byPath := map[string]string{}
	for _, f := range findings {
		byPath[f.Path] = f.Signal
	}
	if byPath["/admin"] != "gray" && byPath["/admin/"] != "gray" {
		t.Fatalf("admin signal=%v", byPath)
	}
	if byPath["/graphql"] != "gray" && byPath["/graphql/"] != "gray" {
		t.Fatalf("graphql signal=%v", byPath)
	}
	if byPath["/swagger"] != "green" && byPath["/swagger/"] != "green" {
		t.Fatalf("swagger signal=%v", byPath)
	}
}

func TestScanReportIncludesAdminSurfaceJSON(t *testing.T) {
	SetPortDialHook(func(host string, port int, _ time.Duration) bool {
		return port == 80
	})
	t.Cleanup(func() { SetPortDialHook(nil) })

	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		switch r.URL.Path {
		case "/":
			w.Write([]byte(` + "`" + r`<html><form enctype="multipart/form-data" action="/api/upload"><input type="file"></form></html>` + "`" + r`))
		case "/api/upload":
			w.WriteHeader(http.StatusOK)
		case "/admin", "/admin/":
			w.WriteHeader(http.StatusForbidden)
		default:
			http.NotFound(w, r)
		}
	}))
	defer srv.Close()

	u, _ := url.Parse(srv.URL)
	port := 80
	if p := u.Port(); p != "" {
		port = atoi(p)
	}

	SetFetchPageHook(func(rawURL string) (int, string, error) {
		resp, err := http.Get(rawURL)
		if err != nil {
			return 0, "", err
		}
		defer resp.Body.Close()
		body, _ := readBodyLimited(resp.Body, maxHTMLBytes)
		return resp.StatusCode, body, nil
	})
	t.Cleanup(func() { SetFetchPageHook(nil) })

	report, err := Scan(ScanRequest{Host: u.Hostname(), Port: port, Scheme: u.Scheme})
	if err != nil {
		t.Fatal(err)
	}
	raw, err := json.Marshal(report)
	if err != nil {
		t.Fatal(err)
	}
	payload := string(raw)
	if !strings.Contains(payload, `"admin_surface"`) {
		t.Fatalf("missing admin_surface: %s", payload)
	}
	if !strings.Contains(payload, `"upload_hunter"`) {
		t.Fatalf("missing upload_hunter: %s", payload)
	}
	if len(report.AdminSurface) == 0 {
		t.Fatal("expected admin surface findings")
	}
	if report.Crawl == nil || len(report.Crawl.UploadHunter) == 0 {
		t.Fatal("expected crawl upload hunter findings")
	}
}
'''


def patch(path: Path, old: str, new: str) -> None:
    text = path.read_text(encoding="utf-8")
    if old not in text:
        if new.strip() in text:
            return
        raise SystemExit(f"patch miss in {path}: {old[:80]!r}")
    path.write_text(text.replace(old, new, 1), encoding="utf-8", newline="\n")


def main() -> None:
    recon = ROOT / "server/internal/recon"
    (recon / "admin_surface.go").write_text(ADMIN_SURFACE, encoding="utf-8", newline="\n")
    (recon / "upload_hunter.go").write_text(UPLOAD_HUNTER, encoding="utf-8", newline="\n")

    types = recon / "types.go"
    patch(
        types,
        "\tCMSFingerprints []string          `json:\"cms_fingerprints,omitempty\"`\n}",
        "\tCMSFingerprints []string              `json:\"cms_fingerprints,omitempty\"`\n"
        "\tUploadHunter    []UploadHunterFinding `json:\"upload_hunter,omitempty\"`\n}",
    )
    patch(
        types,
        "type PageFinding struct {",
        "type UploadHunterFinding struct {\n"
        "\tPageURL    string   `json:\"page_url\"`\n"
        "\tTarget     string   `json:\"target,omitempty\"`\n"
        "\tSource     string   `json:\"source\"`\n"
        "\tMethod     string   `json:\"method,omitempty\"`\n"
        "\tTags       []string `json:\"tags,omitempty\"`\n"
        "\tScore      int      `json:\"score\"`\n"
        "\tStatusCode int      `json:\"status_code,omitempty\"`\n"
        "}\n\n"
        "type AdminSurfaceFinding struct {\n"
        "\tPath       string `json:\"path\"`\n"
        "\tURL        string `json:\"url\"`\n"
        "\tStatusCode int    `json:\"status_code\"`\n"
        "\tSignal     string `json:\"signal\"`\n"
        "}\n\n"
        "type PageFinding struct {",
    )
    patch(
        types,
        "\tCrawl           *CrawlReport           `json:\"crawl,omitempty\"`\n"
        "\tRecommendations []DeployRecommendation `json:\"recommendations,omitempty\"`\n}",
        "\tCrawl           *CrawlReport           `json:\"crawl,omitempty\"`\n"
        "\tAdminSurface    []AdminSurfaceFinding  `json:\"admin_surface,omitempty\"`\n"
        "\tRecommendations []DeployRecommendation `json:\"recommendations,omitempty\"`\n}",
    )

    scan = recon / "scan.go"
    patch(
        scan,
        "\t\tif err == nil && crawl != nil {\n"
        "\t\t\treport.Crawl = crawl\n"
        "\t\t}\n"
        "\t}",
        "\t\tif err == nil && crawl != nil {\n"
        "\t\t\treport.Crawl = crawl\n"
        "\t\t}\n"
        "\t\treport.AdminSurface = ProbeAdminSurface(host, req.Port, req.Scheme)\n"
        "\t}",
    )
    patch(
        scan,
        "\t\tif len(crawl.FileInputs) > 0 || len(crawl.MultipartForms) > 0 {",
        "\t\tif len(crawl.FileInputs) > 0 || len(crawl.MultipartForms) > 0 || len(crawl.UploadHunter) > 0 {",
    )

    crawl = recon / "crawl.go"
    patch(
        crawl,
        "\treport := &CrawlReport{}\n\tvisited := map[string]bool{}",
        "\treport := &CrawlReport{}\n\tvar uploadRaw []UploadHunterFinding\n\tvar jsQueue []string\n\tjsSeen := map[string]bool{}\n\tvisited := map[string]bool{}",
    )
    patch(
        crawl,
        "\t\treport.CMSFingerprints = mergeCMS(report.CMSFingerprints, cms)\n\n\t\tif item.depth >= DefaultCrawlDepth {",
        "\t\treport.CMSFingerprints = mergeCMS(report.CMSFingerprints, cms)\n"
        "\t\tuploadRaw = append(uploadRaw, collectUploadFromPage(item.url, files, multi)...)\n"
        "\t\tuploadRaw = append(uploadRaw, detectDragDropZones(item.url, body)...)\n\n"
        "\t\tif item.depth <= DefaultCrawlDepth {\n"
        "\t\t\tfor _, src := range extractScriptSrc(item.url, body) {\n"
        "\t\t\t\tabs, err := resolveSameOrigin(base, src)\n"
        "\t\t\t\tif err != nil || !sameOrigin(base, abs) || !isScriptAsset(src) {\n"
        "\t\t\t\t\tcontinue\n"
        "\t\t\t\t}\n"
        "\t\t\t\tlkey := normalizeURLKey(abs)\n"
        "\t\t\t\tif jsSeen[lkey] {\n"
        "\t\t\t\t\tcontinue\n"
        "\t\t\t\t}\n"
        "\t\t\t\tjsSeen[lkey] = true\n"
        "\t\t\t\tjsQueue = append(jsQueue, abs)\n"
        "\t\t\t}\n"
        "\t\t}\n\n\t\tif item.depth >= DefaultCrawlDepth {",
    )
    patch(
        crawl,
        "\treport.CMSFingerprints = mergeCMS(nil, report.CMSFingerprints)\n\treturn report, nil",
        "\treport.CMSFingerprints = mergeCMS(nil, report.CMSFingerprints)\n\n"
        "\tmaxJS := 20\n"
        "\tif len(jsQueue) > maxJS {\n"
        "\t\tjsQueue = jsQueue[:maxJS]\n"
        "\t}\n"
        "\tfor _, jsURL := range jsQueue {\n"
        "\t\t_, jsBody, err := fetchPage(jsURL)\n"
        "\t\tif err != nil {\n"
        "\t\t\tcontinue\n"
        "\t\t}\n"
        "\t\tuploadRaw = append(uploadRaw, scanJSForUpload(jsURL, jsBody)...)\n"
        "\t}\n"
        "\treport.UploadHunter = rankUploadFindings(base, uploadRaw)\n\treturn report, nil",
    )

    test = recon / "recon_test.go"
    text = test.read_text(encoding="utf-8")
    if "TestUploadHunterMultipartAndJS" not in text:
        marker = "func containsStr(list []string, want string) bool {"
        if marker not in text:
            raise SystemExit("recon_test.go marker missing")
        text = text.replace(marker, TEST_APPEND + "\n" + marker, 1)
        test.write_text(text, encoding="utf-8", newline="\n")

    print("files written")


def run(cmd: list[str], cwd: Path | None = None) -> None:
    print("+", " ".join(cmd))
    subprocess.run(cmd, cwd=cwd or ROOT, check=True)


if __name__ == "__main__":
    subprocess.run(["git", "checkout", "HEAD", "--", "server/internal/recon/"], cwd=ROOT, check=True)
    main()
    run(["go", "test", "./internal/recon/...", "-count=1"], cwd=ROOT / "server")
    paths = [
        "server/internal/recon/admin_surface.go",
        "server/internal/recon/upload_hunter.go",
        "server/internal/recon/types.go",
        "server/internal/recon/scan.go",
        "server/internal/recon/crawl.go",
        "server/internal/recon/recon_test.go",
    ]
    run(["git", "add", *paths])
    msg = (
        "Add recon upload hunter and admin surface probing for owned-target scans.\n\n"
        "Extend web crawl with multipart/drag-drop/JS upload ranking and probe common "
        "admin paths for 200 vs 401/403 signals in scan JSON."
    )
    run(["git", "commit", "-m", msg])
    run(["git", "pull", "--rebase", "origin", "main"])
    run(["git", "push", "origin", "HEAD"])
    out = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT, text=True).strip()
    print("COMMIT_HASH=" + out)