mirror of
https://github.com/mrvasil/ctf-ad-mirror-tool-rust.git
synced 2026-10-23 16:31:02 +00:00
114 lines
4.3 KiB
Python
114 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
|
"""Disposable loopback-only UI fixture. Never reads legacy configuration.
|
|
|
|
Usage: python3 scripts/ui_fixture.py target/debug/mirror-tool
|
|
Public test-only login: local-ui-fixture-only. Default state is disposable.
|
|
Use --data-dir data/demo for persistent state; --port-offset 100 isolates a test.
|
|
--flag-vault prepares an encrypted synthetic vault; services still default off.
|
|
"""
|
|
import argparse
|
|
import asyncio
|
|
from contextlib import nullcontext
|
|
import os
|
|
from pathlib import Path
|
|
import signal
|
|
import tempfile
|
|
|
|
|
|
async def echo(reader, writer):
|
|
try:
|
|
while data := await reader.read(65536):
|
|
writer.write(data)
|
|
await writer.drain()
|
|
except (OSError, asyncio.CancelledError):
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
try:
|
|
await writer.wait_closed()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
async def main():
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("executable", type=Path)
|
|
parser.add_argument("--flag-vault", action="store_true", help="Configure a synthetic flag vault (service protection remains off)")
|
|
parser.add_argument("--data-dir", type=Path, help="Preserve config/state here; never overwrite an existing config")
|
|
parser.add_argument("--port-offset", type=int, default=0, help="Add to all fixture ports (0..30000)")
|
|
args = parser.parse_args()
|
|
if not 0 <= args.port_offset <= 30000:
|
|
parser.error("port offset must be in 0..30000")
|
|
executable = str(args.executable.resolve())
|
|
ui_port, proxy_port, echo_port = (port + args.port_offset for port in (16700, 18081, 18082))
|
|
if args.data_dir:
|
|
args.data_dir.mkdir(parents=True, exist_ok=True)
|
|
context = nullcontext(args.data_dir.resolve()) if args.data_dir else tempfile.TemporaryDirectory(prefix="mirror-ui-")
|
|
with context as folder:
|
|
root = Path(folder)
|
|
config = root / "mirror.toml"
|
|
new_config = not config.exists()
|
|
if new_config:
|
|
config.write_text(f'''[web]
|
|
listen = "127.0.0.1:{ui_port}"
|
|
state_file = "state.json"
|
|
[network]
|
|
drain_timeout_ms = 500
|
|
[[teams]]
|
|
id = "a"
|
|
name = "Loopback A"
|
|
address = "127.0.0.1"
|
|
[[teams]]
|
|
id = "b"
|
|
name = "Loopback B"
|
|
address = "::1"
|
|
[[services]]
|
|
id = "echo"
|
|
name = "Echo"
|
|
listen = "127.0.0.1:{proxy_port}"
|
|
upstream_port = {echo_port}
|
|
primary = "a"
|
|
replicas = ["b"]
|
|
enabled = false
|
|
auto_switch = false
|
|
replica_policy = "isolate"
|
|
''')
|
|
if args.flag_vault:
|
|
if new_config:
|
|
with config.open("a") as file:
|
|
file.write('\n[flags]\ndatabase = "flags.sqlite3"\n')
|
|
elif "[flags]" not in config.read_text():
|
|
parser.error("existing config has no [flags]; add database = \"flags.sqlite3\" under [flags] first")
|
|
key_file = root / "flag-key.hex"
|
|
if not key_file.exists():
|
|
if (root / "flags.sqlite3").exists():
|
|
parser.error("existing vault has no key file; restore its original key")
|
|
fd = os.open(key_file, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
|
|
with os.fdopen(fd, "w") as file:
|
|
file.write(os.urandom(32).hex())
|
|
servers = [await asyncio.start_server(echo, address, echo_port) for address in ("127.0.0.1", "::1")]
|
|
env = os.environ.copy()
|
|
env["MIRROR_PASSWORD"] = "local-ui-fixture-only"
|
|
if args.flag_vault:
|
|
env["MIRROR_FLAG_KEY"] = key_file.read_text().strip()
|
|
process = await asyncio.create_subprocess_exec(executable, "--config", str(config), env=env)
|
|
stop = asyncio.Event()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
asyncio.get_running_loop().add_signal_handler(sig, stop.set)
|
|
print(f"UI fixture: http://127.0.0.1:{ui_port}; config: {config} (loopback only)", flush=True)
|
|
waiter = asyncio.create_task(process.wait())
|
|
stopper = asyncio.create_task(stop.wait())
|
|
await asyncio.wait([waiter, stopper], return_when=asyncio.FIRST_COMPLETED)
|
|
if process.returncode is None:
|
|
process.terminate()
|
|
await process.wait()
|
|
stopper.cancel()
|
|
await asyncio.gather(stopper, return_exceptions=True)
|
|
for server in servers:
|
|
server.close()
|
|
await server.wait_closed()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|