#!/usr/bin/env python3
import subprocess
from pathlib import Path

ROOT = Path(r"G:/crypto miner")
RECON = ROOT / "server/internal/recon"
SCRIPTS = Path(__file__).parent


def write(name: str, text: str) -> None:
    data = text.encode("utf-8")
    if b"\x00" in data:
        raise ValueError(f"NUL in {name}")
    (RECON / name).write_bytes(data)


def patch_crawl() -> None:
    p = RECON / "crawl.go"
    t = p.read_text(encoding="utf-8")
    if "headerSnaps" not in t:
        t = t.replace("report := &CrawlReport{}\n\tvisited", "report := &CrawlReport{}\n\tvar headerSnaps []HTTPHeaderSnap\n\tvar htmlBodies []string\n\tvisited")
        t = t.replace("status, body, err := fetchPage", "status, body, headers, err := fetchPage")
        t = t.replace(
            "report.Pages = append(report.Pages, PageFinding{URL: item.url, StatusCode: status, Title: title})\n\n\t\tfiles",
            "report.Pages = append(report.Pages, PageFinding{URL: item.url, StatusCode: status, Title: title})\n\t\tif len(headers) > 0 { headerSnaps = append(headerSnaps, HTTPHeaderSnap{URL: item.url, Headers: headers}) }\n\t\thtmlBodies = append(htmlBodies, body)\n\n\t\tfiles",
        )
        t = t.replace(
            "report.CMSFingerprints = mergeCMS(nil, report.CMSFingerprints)\n\treturn report, nil",
            "report.CMSFingerprints = mergeCMS(nil, report.CMSFingerprints)\n\treport.Stack = BuildStack(headerSnaps, htmlBodies)\n\treturn report, nil",
        )
    if "map[string]string, error" not in t:
        t = t.replace("func fetchPage(rawURL string) (int, string, error)", "func fetchPage(rawURL string) (int, string, map[string]string, error)")
        t = t.replace("return fetchPageFn(rawURL)", "s,b,e:=fetchPageFn(rawURL); return s,b,nil,e")
        t = t.replace("return 0, \"\", err", "return 0, \"\", nil, err")
        t = t.replace("return resp.StatusCode, \"\", err", "return resp.StatusCode, \"\", nil, err")
        t = t.replace(
            "return resp.StatusCode, body, nil\n}",
            "headers:=map[string]string{}\n\tfor k,v:=range resp.Header { if len(v)>0 { headers[k]=v[0] } }\n\treturn resp.StatusCode, body, headers, nil\n}",
        )
    p.write_bytes(t.encode("utf-8"))


def patch_tests() -> None:
    p = RECON / "recon_test.go"
    t = p.read_text(encoding="utf-8")
    t = t.replace('ScanPorts("10.0.0.5")', 'ScanPorts("10.0.0.5", nil)')
    t = t.replace('BuildRecommendations(ports, crawl)', 'BuildRecommendations(ports, crawl, nil, "10.0.0.1", false)')
    extra = '''

func TestResolveScanPortsMergesProfiles(t *testing.T) {
	ports, used := ResolveScanPorts([]string{"web", "linux", "cloud_metadata"})
	if !containsInt(ports, 6262) || len(used) != 3 { t.Fatalf("%v %v", ports, used) }
}
func TestBuildStackFromHeaders(t *testing.T) {
	stack := BuildStack([]HTTPHeaderSnap{{Headers: map[string]string{"X-Powered-By":"PHP/8.1"}}}, nil)
	if SuggestDeployKitLane(stack) != "php" { t.Fatal(stack) }
}
func TestGrabBannersWithHooks(t *testing.T) {
	SetBannerHooks(func(_ string,p int) string { if p==22 {return "SSH"}; return "" }, func(_ string,p int)(string,string){ if p==80 {return "t","s"}; return "","" }, func(_ string,p int) string { if p==5985 {return "w"}; return "" }, nil)
	t.Cleanup(func(){SetBannerHooks(nil,nil,nil,nil)})
	if len(GrabBanners("h", []PortResult{{22,true},{80,true},{5985,true}})) != 3 { t.Fatal() }
}
func TestCloudMetadataProfileSuggestsSSM(t *testing.T) {
	SetBannerHooks(nil,nil,nil,func()bool{return true}); t.Cleanup(func(){SetBannerHooks(nil,nil,nil,nil)})
	for _,r := range BuildRecommendations(nil,nil,nil,"ec2.compute.amazonaws.com",true) { if r.Lane=="ssm_document" { return } }
	t.Fatal()
}
func containsInt(a []int,w int) bool { for _,n:=range a { if n==w {return true} }; return false }
'''
    if "TestResolveScanPortsMergesProfiles" not in t:
        t += extra
    p.write_bytes(t.encode("utf-8"))


def patch_recon_ts() -> None:
    p = ROOT / "server/web/src/types/recon.ts"
    t = p.read_text(encoding="utf-8")
    if "profiles_used" in t:
        return
    t = t.replace("paths?: string[];\n}", "paths?: string[];\n  profile?: string;\n  profiles?: string[];\n}")
    t = t.replace(
        "export interface ReconPortResult {\n  port: number;\n  open: boolean;\n}\n\nexport interface ReconFormFinding",
        "export interface ReconPortResult {\n  port: number;\n  open: boolean;\n}\n\nexport interface ReconPortBanner {\n  port: number;\n  service?: string;\n  banner?: string;\n  title?: string;\n  hint?: string;\n}\n\nexport interface ReconStackEntry {\n  name: string;\n  source: string;\n  detail?: string;\n}\n\nexport interface ReconFormFinding",
    )
    t = t.replace("cms_fingerprints?: string[];\n}", "cms_fingerprints?: string[];\n  stack?: ReconStackEntry[];\n}")
    t = t.replace(
        "export interface ReconScanReport {\n  host: string;",
        "export interface ReconScanReport {\n  scan_id?: string;\n  host: string;\n  profile?: string;\n  profiles_used?: string[];\n  status?: string;",
    )
    t = t.replace(
        "ports: ReconPortResult[];\n  crawl?: ReconCrawlReport;",
        "ports: ReconPortResult[];\n  banners?: ReconPortBanner[];\n  stack?: ReconStackEntry[];\n  deploy_kit_lane?: string;\n  crawl?: ReconCrawlReport;",
    )
    p.write_bytes(t.encode("utf-8"))


def patch_subnet_recon() -> None:
    p = ROOT / "agent/deploy/subnet_recon.go"
    t = p.read_text(encoding="utf-8")
    if "SSHBanner" in t:
        return
    t = t.replace("[]int{80, 443, 8080}", "[]int{80, 443, 6262, 8080}")
    t = t.replace(
        'HTTPTitle       string `json:"http_title,omitempty"`\n\tStatus',
        'HTTPTitle       string `json:"http_title,omitempty"`\n\tSSHBanner       string `json:"ssh_banner,omitempty"`\n\tWinRMHint       string `json:"winrm_hint,omitempty"`\n\tStatus',
    )
    t = t.replace(
        "entry.HTTPTitle = title\n\t\t}\n\t\tout = append(out, entry)",
        "entry.HTTPTitle = title\n\t\t}\n\t\tif b:=probeSSHBanner(host,open);b!=\"\"{entry.SSHBanner=b}\n\t\tif h:=probeWinRMHint(host,open);h!=\"\"{entry.WinRMHint=h}\n\t\tout = append(out, entry)",
    )
    insert = '''
func probeSSHBanner(host string, openPorts []int) string {
	for _, p := range openPorts { if p == 22 {
		c, err := net.DialTimeout("tcp", net.JoinHostPort(host,"22"), 2*time.Second)
		if err != nil { return "" }
		defer c.Close()
		b := make([]byte, 256); n, _ := c.Read(b)
		return strings.TrimSpace(string(b[:n]))
	}}
	return ""
}
func probeWinRMHint(host string, openPorts []int) string {
	for _, p := range openPorts { if p == 5985 {
		c, err := net.DialTimeout("tcp", net.JoinHostPort(host,"5985"), 2*time.Second)
		if err == nil { _ = c.Close(); return "winrm_listening" }
	}}
	return ""
}
'''
    t = t.replace("func probeHTTPTitle(host string, openPorts []int) string {", insert + "\nfunc probeHTTPTitle(host string, openPorts []int) string {")
    p.write_bytes(t.encode("utf-8"))


def main() -> None:
    for orphan in ("upload_hunter.go", "admin_surface.go", "recon_ux.go"):
        p = RECON / orphan
        if p.exists():
            p.unlink()
    subprocess.run(["python", str(SCRIPTS / "batch1_build_scan.py")], check=True)
    subprocess.run(["git", "checkout", "HEAD", "--", "server/internal/recon/crawl.go", "server/internal/recon/recon_test.go", "server/internal/recon/relay_scan.go"], cwd=str(ROOT), check=True)
    exec((SCRIPTS / "_batch1_embed.py").read_text(encoding="utf-8"), globals())
    write("types.go", TYPES)
    write("portscan.go", PORTSCAN)
    write("scan.go", (SCRIPTS / "_batch1_scan.go").read_text(encoding="utf-8"))
    ux = RECON / "recon_ux.go"
    if ux.exists():
        ux.unlink()
    for orphan in ("upload_hunter.go", "admin_surface.go", "recon_ux.go"):
        p = RECON / orphan
        if p.exists():
            p.unlink()
    patch_crawl()
    patch_tests()
    patch_recon_ts()
    patch_subnet_recon()
    relay = RECON / "relay_scan.go"
    if relay.exists():
        t = relay.read_text(encoding="utf-8")
        t = t.replace('ScanPorts(host)', 'ScanPorts(host, nil)')
        t = t.replace('BuildRecommendations(ports, nil)', 'BuildRecommendations(ports, nil, nil, host, false)')
        t = t.replace('shell.Recommendations = BuildRecommendations(ports, nil, nil, host, false)', 'shell.Recommendations = BuildRecommendations(ports, nil, nil, shell.Host, false)')
        t = t.replace(
            "Host: r.Host, ScannedAt: r.ScannedAt, Ports: r.Ports, RelayVia: relayVia,\n\t\tUDPHints: r.UDPHints, PathTracerHints: r.PathTracerHints, Message: r.Message,\n\t\tRecommendations: r.Recommendations,",
            "Host: r.Host, ScannedAt: r.ScannedAt, Ports: r.Ports, RelayVia: relayVia, Recommendations: r.Recommendations,",
        )
        relay.write_bytes(t.encode("utf-8"))
    env = {**subprocess.os.environ, "GOCACHE": str(ROOT / ".gocache")}
    r = subprocess.run(["go", "test", "./internal/recon/..."], cwd=str(ROOT / "server"), env=env)
    if r.returncode != 0:
        raise SystemExit(r.returncode)
    subprocess.run(
        ["git", "add", "server/internal/recon", "server/web/src/types/recon.ts", "agent/deploy/subnet_recon.go", "scripts"],
        cwd=str(ROOT), check=True,
    )
    subprocess.run(
        ["git", "commit", "-m", "Add recon network batch 1: stack banners and smart port profiles.\n\nParse technology stack from crawl headers, grab SSH/HTTP/WinRM banners, merge smart port bundles with FleetPorts, and suggest deploy-kit lane plus SSM for EC2 metadata targets."],
        cwd=str(ROOT), check=True,
    )
    subprocess.run(["git", "push", "origin", "main"], cwd=str(ROOT), check=True)
    print("COMMIT_HASH=" + subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=str(ROOT), text=True).strip())


if __name__ == "__main__":
    main()