566 lines
17 KiB
Python
Executable File
566 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import dataclasses
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import random
|
|
import string
|
|
import sys
|
|
from typing import Iterable, List, Optional, Sequence
|
|
|
|
import requests
|
|
|
|
EXIT_OK = 101
|
|
EXIT_CORRUPT = 102
|
|
EXIT_MUMBLE = 103
|
|
EXIT_DOWN = 104
|
|
EXIT_CHECKER_ERROR = 110
|
|
|
|
SCHEME = os.getenv("SIGMA_SCHEME", "http")
|
|
PORT = int(os.getenv("SIGMA_API_PORT") or "3000")
|
|
API_PREFIX = os.getenv("SIGMA_API_PREFIX", "/api/v1").rstrip("/")
|
|
REQUEST_TIMEOUT = float(os.getenv("SIGMA_TIMEOUT", "5"))
|
|
|
|
ALPHABET = string.ascii_lowercase + string.digits
|
|
PASSWORD_ALPHABET = string.ascii_letters + string.digits
|
|
|
|
|
|
def log(message: str) -> None:
|
|
print(message, file=sys.stderr)
|
|
|
|
|
|
class CheckerError(Exception):
|
|
exit_code = EXIT_MUMBLE
|
|
|
|
def __init__(self, message: str):
|
|
super().__init__(message)
|
|
|
|
|
|
class DownError(CheckerError):
|
|
exit_code = EXIT_DOWN
|
|
|
|
|
|
class CorruptError(CheckerError):
|
|
exit_code = EXIT_CORRUPT
|
|
|
|
|
|
class CheckerInternalError(CheckerError):
|
|
exit_code = EXIT_CHECKER_ERROR
|
|
|
|
|
|
class ApiClient:
|
|
def __init__(self, base_url: str, timeout: float = REQUEST_TIMEOUT):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.timeout = timeout
|
|
self.session = requests.Session()
|
|
self.session.headers.update({"User-Agent": "sigma-checker/1.0"})
|
|
|
|
def request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
expected: Optional[Iterable[int]] = None,
|
|
**kwargs,
|
|
) -> requests.Response:
|
|
url = self._build_url(path)
|
|
try:
|
|
response = self.session.request(
|
|
method,
|
|
url,
|
|
timeout=self.timeout,
|
|
**kwargs,
|
|
)
|
|
except requests.RequestException as exc:
|
|
raise DownError(
|
|
f"network error during {method.upper()} {path}: {exc}"
|
|
) from exc
|
|
|
|
if expected is not None:
|
|
allowed = set(expected)
|
|
if response.status_code not in allowed:
|
|
body = response.text[:400]
|
|
raise CheckerError(
|
|
"unexpected status "
|
|
f"{response.status_code} for "
|
|
f"{method.upper()} {path}: {body}"
|
|
)
|
|
|
|
return response
|
|
|
|
def _build_url(self, path: str) -> str:
|
|
if not path.startswith("/"):
|
|
path = "/" + path
|
|
return f"{self.base_url}{path}"
|
|
|
|
|
|
def seeded_random(flag_id: str, vuln: int) -> random.Random:
|
|
digest = hashlib.sha256(f"{flag_id}:{vuln}".encode()).digest()
|
|
return random.Random(int.from_bytes(digest, "big"))
|
|
|
|
|
|
def random_string(
|
|
rng: random.Random,
|
|
length: int,
|
|
alphabet: str = ALPHABET,
|
|
) -> str:
|
|
return "".join(rng.choice(alphabet) for _ in range(length))
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class UserSpec:
|
|
username: str
|
|
password: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class ProjectSpec:
|
|
name: str
|
|
slug: str
|
|
description_template: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class FunctionSpec:
|
|
name: str
|
|
path: str
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Vuln1Context:
|
|
user: UserSpec
|
|
primary_project: ProjectSpec
|
|
decoy_project: ProjectSpec
|
|
|
|
|
|
@dataclasses.dataclass
|
|
class Vuln2Context:
|
|
user: UserSpec
|
|
project: ProjectSpec
|
|
store_function: FunctionSpec
|
|
fetch_function: FunctionSpec
|
|
|
|
|
|
def build_vuln1_context(flag_id: str) -> Vuln1Context:
|
|
rng = seeded_random(flag_id, 1)
|
|
username = f"user_{random_string(rng, 10)}"
|
|
password = random_string(rng, 16, PASSWORD_ALPHABET)
|
|
primary = ProjectSpec(
|
|
name=f"Sigma {random_string(rng, 5).title()}",
|
|
slug=f"proj{random_string(rng, 10)}",
|
|
description_template="Primary telemetry secret: {flag}",
|
|
)
|
|
decoy = ProjectSpec(
|
|
name=f"Telemetry {random_string(rng, 5).title()}",
|
|
slug=f"proj{random_string(rng, 10)}",
|
|
description_template="Routine system baseline",
|
|
)
|
|
return Vuln1Context(UserSpec(username, password), primary, decoy)
|
|
|
|
|
|
def build_vuln2_context(flag_id: str) -> Vuln2Context:
|
|
rng = seeded_random(flag_id, 2)
|
|
username = f"{random_string(rng, 12)}"
|
|
password = random_string(rng, 18, PASSWORD_ALPHABET)
|
|
project = ProjectSpec(
|
|
name=f"Runner {random_string(rng, 6).title()}",
|
|
slug=f"proj{random_string(rng, 12)}",
|
|
description_template="Serverless SQL vault",
|
|
)
|
|
store_fn = FunctionSpec(
|
|
name=f"store_{random_string(rng, 6)}",
|
|
path=f"/{random_string(rng, 6)}",
|
|
)
|
|
fetch_fn = FunctionSpec(
|
|
name=f"fetch_{random_string(rng, 6)}",
|
|
path=f"/{random_string(rng, 6)}",
|
|
)
|
|
return Vuln2Context(
|
|
UserSpec(username, password),
|
|
project,
|
|
store_fn,
|
|
fetch_fn,
|
|
)
|
|
|
|
|
|
class SigmaChecker:
|
|
def __init__(
|
|
self,
|
|
host: str,
|
|
*,
|
|
scheme: str = SCHEME,
|
|
port: Optional[int] = PORT,
|
|
) -> None:
|
|
base_url = self._build_base_url(host, scheme, port)
|
|
prefix = API_PREFIX if API_PREFIX else ""
|
|
self.client = ApiClient(f"{base_url}{prefix}")
|
|
|
|
@staticmethod
|
|
def _build_base_url(host: str, scheme: str, port: Optional[int]) -> str:
|
|
hostname = host
|
|
if ":" in host and not host.startswith("["):
|
|
hostname = f"[{host}]"
|
|
if port is None:
|
|
return f"{scheme}://{hostname}"
|
|
return f"{scheme}://{hostname}:{port}"
|
|
|
|
def check(self) -> None:
|
|
response = self.client.request("GET", "/health/cpu", expected={200})
|
|
try:
|
|
payload = response.json()
|
|
except ValueError as exc:
|
|
raise CheckerError(
|
|
"health endpoint returned invalid JSON"
|
|
) from exc
|
|
usage = payload.get("usage")
|
|
if not isinstance(usage, (int, float)):
|
|
raise CheckerError("health payload missing numeric usage")
|
|
|
|
def put(self, flag_id: str, flag: str, vuln: int) -> None:
|
|
if vuln == 1:
|
|
self._put_vuln1(build_vuln1_context(flag_id), flag)
|
|
elif vuln == 2:
|
|
self._put_vuln2(build_vuln2_context(flag_id), flag)
|
|
else:
|
|
raise CheckerInternalError(f"unsupported vuln {vuln}")
|
|
|
|
def get(self, flag_id: str, flag: str, vuln: int) -> None:
|
|
if vuln == 1:
|
|
self._get_vuln1(build_vuln1_context(flag_id), flag)
|
|
elif vuln == 2:
|
|
self._get_vuln2(build_vuln2_context(flag_id), flag)
|
|
else:
|
|
raise CheckerInternalError(f"unsupported vuln {vuln}")
|
|
|
|
def _ensure_user(self, user: UserSpec) -> None:
|
|
payload = {"username": user.username, "password": user.password}
|
|
response = self.client.request(
|
|
"POST",
|
|
"/users",
|
|
json=payload,
|
|
expected={201, 409},
|
|
)
|
|
if response.status_code == 201:
|
|
log(f"registered {user.username}")
|
|
|
|
def _login(self, user: UserSpec, *, corrupt_on_fail: bool = False) -> str:
|
|
payload = {"username": user.username, "password": user.password}
|
|
response = self.client.request(
|
|
"POST",
|
|
"/auth/login",
|
|
json=payload,
|
|
expected={200, 401},
|
|
)
|
|
if response.status_code == 401:
|
|
if corrupt_on_fail:
|
|
raise CorruptError("stored credentials rejected")
|
|
raise CheckerError("authentication failed")
|
|
try:
|
|
body = response.json()
|
|
except ValueError as exc:
|
|
raise CheckerError("login endpoint returned invalid JSON") from exc
|
|
token = body.get("auth_token")
|
|
if not isinstance(token, str) or not token:
|
|
raise CheckerError("login response missing token")
|
|
return token
|
|
|
|
def _auth_headers(self, token: str) -> dict:
|
|
return {
|
|
"Authorization": f"Bearer {token}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
def _create_project(
|
|
self,
|
|
headers: dict,
|
|
spec: ProjectSpec,
|
|
flag: Optional[str] = None,
|
|
) -> dict:
|
|
description = spec.description_template
|
|
if flag is not None:
|
|
description = description.format(flag=flag)
|
|
payload = {
|
|
"name": spec.name,
|
|
"slug": spec.slug,
|
|
"description": description,
|
|
}
|
|
response = self.client.request(
|
|
"POST",
|
|
"/projects",
|
|
headers=headers,
|
|
json=payload,
|
|
expected={201},
|
|
)
|
|
try:
|
|
project = response.json()
|
|
except ValueError as exc:
|
|
raise CheckerError(
|
|
"project creation returned invalid JSON"
|
|
) from exc
|
|
if not isinstance(project.get("_id"), str):
|
|
raise CheckerError("project creation missing id")
|
|
return project
|
|
|
|
def _list_projects(self, headers: dict) -> List[dict]:
|
|
response = self.client.request(
|
|
"GET",
|
|
"/projects",
|
|
headers=headers,
|
|
expected={200},
|
|
)
|
|
try:
|
|
projects = response.json()
|
|
except ValueError as exc:
|
|
raise CheckerError("projects list returned invalid JSON") from exc
|
|
if not isinstance(projects, list):
|
|
raise CheckerError("projects response is not a list")
|
|
return projects
|
|
|
|
def _find_project_by_slug(
|
|
self,
|
|
headers: dict,
|
|
slug: str,
|
|
) -> Optional[dict]:
|
|
for project in self._list_projects(headers):
|
|
if project.get("slug") == slug:
|
|
return project
|
|
return None
|
|
|
|
def _create_function(
|
|
self,
|
|
headers: dict,
|
|
project_id: str,
|
|
spec: FunctionSpec,
|
|
) -> dict:
|
|
payload = {"name": spec.name, "path": spec.path}
|
|
response = self.client.request(
|
|
"POST",
|
|
f"/projects/{project_id}/functions",
|
|
headers=headers,
|
|
json=payload,
|
|
expected={201},
|
|
)
|
|
try:
|
|
func = response.json()
|
|
except ValueError as exc:
|
|
raise CheckerError(
|
|
"function creation returned invalid JSON"
|
|
) from exc
|
|
if not isinstance(func.get("_id"), str):
|
|
raise CheckerError("function creation missing id")
|
|
return func
|
|
|
|
def _update_function(
|
|
self,
|
|
headers: dict,
|
|
project_id: str,
|
|
function_id: str,
|
|
code: str,
|
|
methods: Sequence[str],
|
|
) -> None:
|
|
payload = {"code": code, "methods": list(methods)}
|
|
self.client.request(
|
|
"PUT",
|
|
f"/projects/{project_id}/functions/{function_id}",
|
|
headers=headers,
|
|
json=payload,
|
|
expected={200},
|
|
)
|
|
|
|
def _execute_function(
|
|
self,
|
|
method: str,
|
|
slug: str,
|
|
path: str,
|
|
expected: Iterable[int] = (200,),
|
|
) -> requests.Response:
|
|
return self.client.request(
|
|
method,
|
|
f"/exec/{slug}{path}",
|
|
expected=set(expected),
|
|
)
|
|
|
|
def _put_vuln1(self, context: Vuln1Context, flag: str) -> None:
|
|
self._ensure_user(context.user)
|
|
token = self._login(context.user)
|
|
headers = self._auth_headers(token)
|
|
self._create_project(headers, context.primary_project, flag)
|
|
self._create_project(headers, context.decoy_project)
|
|
stored = self._find_project_by_slug(
|
|
headers,
|
|
context.primary_project.slug,
|
|
)
|
|
if not stored:
|
|
raise CheckerError("project missing after creation")
|
|
description = stored.get("description", "")
|
|
if flag not in description:
|
|
raise CheckerError("flag not present immediately after storing")
|
|
|
|
def _get_vuln1(self, context: Vuln1Context, flag: str) -> None:
|
|
token = self._login(context.user, corrupt_on_fail=True)
|
|
headers = self._auth_headers(token)
|
|
project = self._find_project_by_slug(
|
|
headers,
|
|
context.primary_project.slug,
|
|
)
|
|
if not project:
|
|
raise CorruptError("flag project missing")
|
|
description = project.get("description", "")
|
|
if flag not in description:
|
|
raise CorruptError("flag missing from project description")
|
|
|
|
def _put_vuln2(self, context: Vuln2Context, flag: str) -> None:
|
|
self._ensure_user(context.user)
|
|
token = self._login(context.user)
|
|
headers = self._auth_headers(token)
|
|
project = self._create_project(headers, context.project)
|
|
project_id = project["_id"]
|
|
func = self._create_function(
|
|
headers,
|
|
project_id,
|
|
context.store_function,
|
|
)
|
|
code = self._build_store_function_code(context.project.slug, flag)
|
|
self._update_function(headers, project_id, func["_id"], code, ["POST"])
|
|
self._execute_function(
|
|
"POST",
|
|
context.project.slug,
|
|
context.store_function.path,
|
|
expected=(200, 201),
|
|
)
|
|
|
|
def _get_vuln2(self, context: Vuln2Context, flag: str) -> None:
|
|
token = self._login(context.user, corrupt_on_fail=True)
|
|
headers = self._auth_headers(token)
|
|
project = self._find_project_by_slug(headers, context.project.slug)
|
|
if not project:
|
|
raise CorruptError("project not found during GET")
|
|
project_id = project.get("_id")
|
|
if not isinstance(project_id, str):
|
|
raise CheckerError("project missing id")
|
|
func = self._create_function(
|
|
headers,
|
|
project_id,
|
|
context.fetch_function,
|
|
)
|
|
code = self._build_fetch_function_code(context.project.slug)
|
|
self._update_function(headers, project_id, func["_id"], code, ["GET"])
|
|
response = self._execute_function(
|
|
"GET",
|
|
context.project.slug,
|
|
context.fetch_function.path,
|
|
expected=(200, 404),
|
|
)
|
|
if response.status_code == 404:
|
|
raise CorruptError("flag table returned empty result")
|
|
payload = self._parse_json(response.text)
|
|
if payload.get("flag") != flag:
|
|
raise CorruptError("retrieved flag mismatch")
|
|
|
|
def _build_store_function_code(self, slug: str, flag: str) -> str:
|
|
create_sql = (
|
|
f'CREATE TABLE IF NOT EXISTS "scope_{slug}" '
|
|
"(id SERIAL PRIMARY KEY, flag TEXT NOT NULL)"
|
|
)
|
|
insert_sql = f'INSERT INTO "scope_{slug}" (flag) VALUES ($1)'
|
|
return (
|
|
"""
|
|
;(async () => {
|
|
const CREATE_SQL = __CREATE__;
|
|
const INSERT_SQL = __INSERT__;
|
|
const FLAG_VALUE = __FLAG__;
|
|
await sql.execute(CREATE_SQL);
|
|
await sql.execute(INSERT_SQL, [FLAG_VALUE]);
|
|
res.setStatus(201).json({ status: "stored" });
|
|
})().catch((err) => {
|
|
res.setStatus(500).json({ error: String(err?.message ?? err) });
|
|
});
|
|
"""
|
|
.replace("__CREATE__", json.dumps(create_sql))
|
|
.replace("__INSERT__", json.dumps(insert_sql))
|
|
.replace("__FLAG__", json.dumps(flag))
|
|
.strip()
|
|
)
|
|
|
|
def _build_fetch_function_code(self, slug: str) -> str:
|
|
select_sql = f'SELECT flag FROM "scope_{slug}" ORDER BY id DESC LIMIT 1'
|
|
return (
|
|
"""
|
|
;(async () => {
|
|
const SELECT_SQL = __SELECT__;
|
|
const rows = await sql.query(SELECT_SQL);
|
|
if (
|
|
!Array.isArray(rows) ||
|
|
rows.length === 0 ||
|
|
typeof rows[0]?.flag !== "string"
|
|
) {
|
|
res.setStatus(404).json({ error: "flag missing" });
|
|
return;
|
|
}
|
|
res.setStatus(200).json({ flag: rows[0].flag });
|
|
})().catch((err) => {
|
|
res.setStatus(500).json({ error: String(err?.message ?? err) });
|
|
});
|
|
"""
|
|
.replace("__SELECT__", json.dumps(select_sql))
|
|
.strip()
|
|
)
|
|
|
|
@staticmethod
|
|
def _parse_json(body: str) -> dict:
|
|
try:
|
|
data = json.loads(body or "null")
|
|
except json.JSONDecodeError as exc:
|
|
raise CheckerError("function response is not valid JSON") from exc
|
|
if not isinstance(data, dict):
|
|
raise CheckerError("function response JSON is not an object")
|
|
return data
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Sigma ForcAD checker")
|
|
parser.add_argument(
|
|
"action",
|
|
choices=["check", "put", "get"],
|
|
help="Checker action",
|
|
)
|
|
parser.add_argument("host", help="Team IP address")
|
|
parser.add_argument("flag_id", nargs="?")
|
|
parser.add_argument("flag", nargs="?")
|
|
parser.add_argument("vuln", nargs="?", type=int)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
checker = SigmaChecker(args.host)
|
|
try:
|
|
if args.action == "check":
|
|
checker.check()
|
|
elif args.action == "put":
|
|
if not (args.flag_id and args.flag and args.vuln is not None):
|
|
raise CheckerInternalError("put requires flag_id flag vuln")
|
|
checker.put(args.flag_id, args.flag, args.vuln)
|
|
elif args.action == "get":
|
|
if not (args.flag_id and args.flag and args.vuln is not None):
|
|
raise CheckerInternalError("get requires flag_id flag vuln")
|
|
checker.get(args.flag_id, args.flag, args.vuln)
|
|
else:
|
|
raise CheckerInternalError("unknown action")
|
|
except CheckerError as exc:
|
|
log(str(exc))
|
|
return exc.exit_code
|
|
except Exception as exc: # pragma: no cover
|
|
log(f"unexpected checker error: {exc}")
|
|
return EXIT_CHECKER_ERROR
|
|
log("OK")
|
|
return EXIT_OK
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|