#!/usr/bin/env python3 import hashlib import sys import secrets import json import checklib import requests import random import websockets import websockets.sync.client from requests.cookies import get_cookie_header import level_gen def random_str(prefix="") -> str: if random.random() < 0.1: return prefix + "Hello world" elif random.random() < 0.1: return prefix + "How are you?" elif random.random() < 0.1: return prefix + "I'm fine" elif random.random() < 0.1: return prefix + "Super secret communication" elif random.random() < 0.1: return prefix + "There are flags I want to send you" else: return prefix + checklib.rnd_string(5 + secrets.randbelow(40)) def random_username() -> str: if random.random() < 0.8: return checklib.rnd_username() return checklib.rnd_string(5 + secrets.randbelow(20)) def random_password() -> str: if random.random() < 0.8: return checklib.rnd_password() return checklib.rnd_string(5 + secrets.randbelow(20)) def random_level_name(hint = "") -> str: x = random.random() if x < 0.33: return checklib.rnd_string(5 + secrets.randbelow(20)) + "_" + hint elif x < 0.66: return checklib.rnd_string(5 + secrets.randbelow(20)) + "_" + hint + "_" + checklib.rnd_string(5 + secrets.randbelow(20)) else: return hint + "_" + checklib.rnd_string(5 + secrets.randbelow(20)) def random_description(hint = "") -> str: strings = [checklib.rnd_string(5 + secrets.randbelow(20)) for _ in range(random.randrange(1, 3))] strings.insert(random.randrange(len(strings)), hint) return "\n".join(strings) def random_prize(hint = "") -> str: x = random.random() if x < 0.33: return "youre " + checklib.rnd_string(5 + secrets.randbelow(20)) + " cool" if x < 0.66: return "secret secret " +checklib.rnd_string(5 + secrets.randbelow(20)) return "my promocode: " + checklib.rnd_string(5 + secrets.randbelow(20)) def minidumps(data) -> str: return json.dumps(data, separators=(",", ":"), sort_keys=True) def jsonhash(data: dict) -> str: return hashlib.sha256(minidumps(data).encode()).hexdigest() class Checker(checklib.BaseChecker): vulns = 2 timeout = 60 uses_attack_data: bool = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) port = 5874 self.host_with_port = f"{self.host}:{port}" self.url = f"http://{self.host_with_port}" def action(self, action, *args, **kwargs): try: super().action(action, *args, **kwargs) except self.get_check_finished_exception(): raise except requests.RequestException as e: self.cquit( checklib.Status.DOWN, "Connection error", f"Requests error {repr(e)}" ) def check_whoami(self, sess: requests.Session, username: str): r = sess.get(f"{self.url}/api/user") self.assert_( r.ok, f"Could not get whoami after registering: {r.status_code=} {r.text=}" ) data = self.get_json(r, f"Invalid json in whoami response: {r.text=}") self.assert_("id" in data, f"Invalid whoami response: no id, {r.text=}") self.assert_( "username" in data, f"Invalid whoami response: no username, {r.text=}" ) self.assert_( data["username"] == username, f"Invalid whoami response: username mismatch, {r.text=}", ) def random_new_user(self) -> tuple[str, str, requests.Session]: sess = self.get_initialized_session() username = random_username() password = random_password() r = sess.post( f"{self.url}/api/auth/register", json={ "username": username, "password": password, }, ) self.assert_(r.ok, f"Could not register: {r.status_code=} {r.text=}") self.check_whoami(sess, username) return username, password, sess def login( self, username: str, password: str, check_whoami: bool = True, error_status: checklib.Status = checklib.Status.MUMBLE, ) -> requests.Session: sess = self.get_initialized_session() r = sess.post( f"{self.url}/api/auth/login", json={ "username": username, "password": password, }, ) self.assert_(r.ok, f"Could not login: {r.status_code=} {r.text=}", error_status) if check_whoami: self.check_whoami(sess, username) return sess def get_level_list( self, sess: requests.Session, checker_status: checklib.Status = checklib.Status.MUMBLE, ) -> list: r = sess.get(f"{self.url}/api/user/levels?page=0") self.assert_( r.ok, f"Could not get level list: {r.status_code=} {r.text=}", checker_status, ) data = self.get_json( r, f"Invalid json in level list: {r.text=}", checker_status ) self.assert_( isinstance(data, list), f"Invalid level list: not a list, {r.text=}", checker_status, ) return data def verify_level(self, level: dict): self.assert_( isinstance(level, dict), f"Level data {level} is invalid: expected dict" ) self.assert_( "id" in level and isinstance(level["id"], int), f"Level data {level} is invalid: expected id of type int", ) self.assert_( "name" in level and isinstance(level["name"], str), f"Level data {level} is invalid: expected name of type str", ) self.assert_( "description" in level and isinstance(level["description"], str), f"Level data {level} is invalid: expected description of type str", ) self.assert_( "visibility" in level and level["visibility"] in ("private", "public"), f"Level data {level} is invalid: expected visibility 'private' or 'public'", ) return True def get_level_by_name( self, sess: requests.Session, name: str, not_found_status: checklib.Status = checklib.Status.MUMBLE, ) -> dict: r = sess.get(f"{self.url}/api/user/level", params={"name": name}) if r.status_code == 404: self.assert_( False, f"Level {name=} not found via search", not_found_status, ) self.assert_( r.ok, f"Could not search level by name: {r.status_code=} {r.text=}", not_found_status, ) level_data = self.get_json( r, f"Invalid json in level search response: {r.text=}", not_found_status ) self.verify_level(level_data) self.assert_eq(level_data["name"], name, f"{level_data['name']=} mismatch") return level_data def check_level_saved( self, sess: requests.Session, name: str, description: str, visibility: str, data: dict, prize: str = None, ): level_list_data = self.get_level_by_name(sess, name) level_id = level_list_data["id"] self.assert_eq( level_list_data["name"], name, f"{level_list_data['name']=} mismatch" ) self.assert_eq( level_list_data["description"], description, f"{level_list_data['description']=} mismatch", ) self.assert_eq( level_list_data["visibility"], visibility, f"{level_list_data['visibility']=} mismatch", ) r = sess.get(f"{self.url}/api/user/level/{level_id}") level_data = self.get_json(r, f"Invalid json in level data: {r.text}") self.verify_level(level_data) self.assert_( "data" in level_data, f"Level data {level_data} is invalid: expected level data", ) self.assert_eq(level_data["name"], name, f"{level_data['name']=} mismatch") self.assert_eq( level_data["description"], description, f"{level_data['description']=} mismatch", ) self.assert_eq( level_data["visibility"], visibility, f"{level_data['visibility']=} mismatch", ) self.assert_eq(level_data["data"], data, f"{level_data['data']=} mismatch") if prize is not None: self.assert_( "prize" in level_data and isinstance(level_data["prize"], str), f"Level data {level_data} is invalid: expected prize of type string", ) self.assert_eq( level_data["prize"], prize, f"{level_data['data']=} mismatch" ) def put(self, flag_id: str, flag: str, vuln: str): username, password, sess = self.random_new_user() try: vuln = int(vuln) except ValueError as e: raise ValueError(f"Invalid {vuln=}, expected a number") from e if vuln == 1: tiles_data = level_gen.generate_unbeatable() level_name = f"unbeatable_FLAG_by_{username}_{secrets.token_hex(4)}" description = "Win to get the FLAG. This level is impossible to beat!!" visibility = "public" elif vuln == 2: tiles_data, _ = level_gen.generate_easy() level_name = f"private_FLAG_by_{username}_{secrets.token_hex(4)}" description = "Win to get the FLAG" visibility = "private" else: raise ValueError(f"Invalid {vuln=}, expected 1 or 2") r = sess.post( f"{self.url}/api/user/level", json={ "name": level_name, "description": description, "visibility": visibility, "data": tiles_data, "prize": flag, }, ) self.assert_(r.ok, f"Could not create the level: {r.status_code=} {r.text=}") self.check_level_saved( sess, level_name, description, visibility, tiles_data, flag ) tiles_data_hash = jsonhash(tiles_data) self.cquit( checklib.Status.OK, public=minidumps({"level_name": level_name}), private=minidumps((username, password, level_name, visibility, tiles_data_hash)), ) def get(self, flag_id: str, flag: str, vuln: str): username, password, level_name, visibility, tiles_data_hash = json.loads(flag_id) sess = self.login( username, password, check_whoami=False, error_status=checklib.Status.CORRUPT ) level_list_data = self.get_level_by_name( sess, level_name, not_found_status=checklib.Status.CORRUPT ) r = sess.get(f"{self.url}/api/user/level/{level_list_data['id']}") level_data = self.get_json(r, f"Invalid json in level data: {r.text}") self.verify_level(level_data) self.assert_( "prize" in level_data and isinstance(level_data["prize"], str), f"Level data {level_data} is invalid: expected prize of type string", ) self.assert_eq( level_data["prize"], flag, "Flag mismatch", checklib.Status.CORRUPT ) self.assert_( "data" in level_data, f"Level data {level_data} is invalid: expected tiles data", ) actual_tiles_data_hash = jsonhash(level_data["data"]) self.assert_eq(tiles_data_hash, actual_tiles_data_hash, "Flag level contents changed", checklib.Status.CORRUPT) self.assert_eq(visibility, level_data["visibility"], "Flag level visibility changed", checklib.Status.CORRUPT) self.cquit(checklib.Status.OK) def check_user_search(self, sess, username): r = sess.post(f"{self.url}/api/user", json={"username": username}) data = self.get_json(r, f"Invalid POST /api/user response json: {r.text}") self.assert_in("username", data, f"Expected username in POST /api/user response: {r.text}") self.assert_eq(data["username"], username, f"Expected {username=}: {r.text}") self.assert_in("id", data, f"Expected id in POST /api/user response: {r.text}") id = data["id"] r = sess.post(f"{self.url}/api/user", json={"id": id}) data = self.get_json(r, f"Invalid POST /api/user response json: {r.text}") self.assert_in("username", data, f"Expected username in POST /api/user response: {r.text}") self.assert_eq(data["username"], username, f"Expected {username=}: {r.text}") self.assert_in("id", data, f"Expected id in POST /api/user response: {r.text}") self.assert_eq(data["id"], id, f"Expected {id=}: {r.text}") def get_level_data(self, sess, id): r = sess.get(f"{self.url}/api/user/level/{id}") level_data = self.get_json(r, f"Invalid json in level data: {r.text}") self.verify_level(level_data) return level_data def connect_with_auth(self, sess: requests.Session, url): cookie_value = get_cookie_header(sess.cookies, requests.Request("GET", url.replace("ws://", "http://"))) return websockets.sync.client.connect(url, additional_headers={"Cookie": cookie_value}) def recv_message(self, sock, timeout): try: m = sock.recv(timeout) except TimeoutError: return False except websockets.exceptions.ConnectionClosedError: return False try: m = json.loads(m) except json.JSONDecodeError: self.assert_(False, f"Invalid websocket message: {m}") if "type" in m and m["type"] == "level_complete": self.assert_in("option", m, f"Invalid level_complete: {m}") self.assert_in("prize", m["option"], f"Invalid level_complete: {m}") return m["option"]["prize"] return True def play_level(self, sess, id, moves = [], expect_prize = False, expect_updates = False): got_updates = False with self.connect_with_auth(sess, f"ws://{self.host_with_port}/api/user/level/{id}/play") as sock: for move in moves: sock.send(json.dumps({"type": "move", "option": {"direction": move}})) while True: timeout = 0.5 if expect_updates and not got_updates else 0.01 m = self.recv_message(sock, timeout) if m == True: # noqa: E712 got_updates = True elif m == False: # noqa: E712 break else: return m while True: timeout = 0.5 if expect_prize else 0.05 m = self.recv_message(sock, timeout) if m == True: # noqa: E712 got_updates = True elif m == False: # noqa: E712 break else: return m self.assert_(not expect_prize, "Didn't get the prize") self.assert_(got_updates or not expect_updates, "Didn't get any updates from the server") def check(self): username, password, sess = self.random_new_user() self.check_user_search(sess, username) r = sess.post(f"{self.url}/api/auth/logout") self.assert_(r.ok, f"Couldn't logout: {r.text}") r = sess.post(f"{self.url}/api/user/whoami") self.assert_(not r.ok, f"Can access whoami after logout: {r.text}") sess = self.login(username, password) levels = self.get_level_list(sess) if len(levels) > 0: level = random.choice(levels) data = self.get_level_data(sess, level["id"]) self.play_level(sess, data["id"]) if random.random() < 0.5: tiles_data, moves = level_gen.generate_easy() level_name = random_level_name("ez") description = random_description("Easy level") else: tiles_data, moves = level_gen.generate_box_test() level_name = random_level_name("box") description = random_description("boxes boxes") prize = random_prize() r = sess.post( f"{self.url}/api/user/level", json={ "name": level_name, "description": description, "visibility": "public", "data": tiles_data, "prize": prize, }, ) data = self.get_json(r, f"Invalid level creation response: {r.text}") self.assert_in("id", data, f"Expected id after level creation: {r.text}") id = data["id"] got_prize = self.play_level(sess, id, moves, True, True) self.assert_eq(prize, got_prize, f"Got wrong prize {got_prize}") self.check_level_saved(sess, level_name, description, "public", tiles_data, prize) username, password, sess = self.random_new_user() got_prize = self.play_level(sess, id, moves, True, True) self.assert_eq(prize, got_prize, f"Got wrong prize {got_prize}") self.check_level_saved(sess, level_name, description, "public", tiles_data) tiles_data, moves = level_gen.generate_box_test() level_name = random_level_name("boxes_priv") description = random_description("only mine level") prize = random_prize() r = sess.post( f"{self.url}/api/user/level", json={ "name": level_name, "description": description, "visibility": "private", "data": tiles_data, "prize": prize, }, ) data = self.get_json(r, f"Invalid level creation response: {r.text}") self.assert_in("id", data, f"Expected id after level creation: {r.text}") id = data["id"] self.check_level_saved(sess, level_name, description, "private", tiles_data, prize) username, password, sess = self.random_new_user() r = sess.get(f"{self.url}/api/user/level/{id}") self.assert_(not r.ok, f"Successfully got private level data: {id} by {username}") r = sess.get(f"{self.url}/api/user/level/{id}/play") self.assert_(not r.ok, f"Successfully started playing private level: {id} played by {username}") self.cquit(checklib.Status.OK) if __name__ == "__main__": host = sys.argv[2] checker = Checker(host) try: action = sys.argv[1] arguments = sys.argv[3:] checker.action(action, *arguments) except checker.get_check_finished_exception(): checklib.cquit( checklib.Status(checker.status), checker.public, checker.private, )