#!/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 from functools import lru_cache from time import sleep #from pwn import remote, p8, u8, u16, p16, flat, pack, unpack, context, args from pwnlib.util.packing import p8, u8, u16, p16, flat, unpack, pack from pwnlib.tubes.remote import remote from Crypto.Cipher import ARC4 from grob_enums import ( Cmd, Status, MovePayload, Color, Piece, CmdGame ) import chess import chess.pgn EXIT_OK = 101 EXIT_CORRUPT = 102 EXIT_MUMBLE = 103 EXIT_DOWN = 104 EXIT_CHECKER_ERROR = 110 PORT = int(os.getenv("GROB_PORT") or "11331") TIMEOUT = float(os.getenv("GROB_TIMEOUT", "1")) ALPHABET = string.ascii_uppercase + 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 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 Comment: color: int after_move: int @dataclasses.dataclass class VulnContext: pgn_id: int comm: Comment room_code: str room_pass: str @lru_cache def _avail_games(): return len(os.listdir('/checkers/grob/games')) def _load_pgn(pgn_idx): with open(f'/checkers/grob/games/{pgn_idx:04d}.pgn', 'r') as pgn_file: return chess.pgn.read_game(pgn_file) def build_vuln_context(flag_id: str) -> VulnContext: rng = seeded_random(flag_id, 1) pgn_id = rng.randint(1, _avail_games()) pgn = _load_pgn(pgn_id) moves = pgn.end().ply() comm = Comment(rng.randint(0, 1), rng.randint(0, moves - 2)) room_code = random_string(rng, 16).encode() room_pass = random_string(rng, 16).encode() return VulnContext(pgn_id, comm, room_code, room_pass) def u128(n, *args, **kwargs): return unpack(n, *args, word_size=128, **kwargs) def p128(n, *args, **kwargs): return pack(n, *args, word_size=128, **kwargs) class arc4tube(remote): MODULO = (1 << 128) - 159 def __init__(self, *a, **kw): super().__init__(*a, **kw) priv = u128(random.randbytes(16)) pub = pow(5, priv, self.MODULO) server_pub = u128(self.recvn(16)) self.send(p128(pub)) cph = ARC4.new(p128(pow(server_pub, priv, self.MODULO))) self.recv_raw = lambda *a, **kw: cph.encrypt(data) if (data := super(arc4tube, self).recv_raw(*a, **kw)) else b'' self.send_raw = lambda data, *a, **kw: super(arc4tube, self).send_raw(data and cph.encrypt(data), *a, **kw) class GrobChecker: def __init__(self, host: str, *, port: Optional[int] = PORT) -> None: self.host = host self.port = port self._is_prev_line_msg = 0 def check(self) -> None: rng = random.Random() pgn_id = rng.randint(1, _avail_games()) moves = _load_pgn(pgn_id).end().ply() send_msg_color = rng.randint(0, 1) send_msg_after_move = rng.randint(0, moves - 2) msg = 'WoW! It\'s very interesting...' room_code = random_string(rng, 16).encode() room_pass = random_string(rng, 16).encode() valid_pgn = bytearray(b'[Password "' + room_pass + b'"]\n') with self._connect() as io1, self._connect() as io2: self._enter_room(room_code, room_pass, io1, io2) color = u8(io1.recvn(1)) io2.recvn(1) if not color: io1, io2 = io2, io1 game_end = None order = True for i, move in enumerate(self._pgn_iter(pgn_id, valid_pgn)): if game_end is not None and not isinstance(move, str): raise CheckerError('game engine fault: invalid game end') if isinstance(move, str): if game_end != move: raise CheckerError('game engine fault: invalid game end') return if order: m = flat(CmdGame.CMD_GAME_MAKE_MOVE, *move, word_size=8) io1.send(m) s = u8(io1.recvn(1)) if s not in (Status.STATUS_SUCCESS, Status.STATUS_GAME_CHECKMATE, Status.STATUS_GAME_DRAW): raise CheckerError('game engine fault: invalid move verdict') if not (io2.recv(1) != s and io2.recvn(6) != m): raise CheckerError('game engine fault: opponent got invalid move') if s == Status.STATUS_GAME_CHECKMATE: game_end = 'W' if order else 'B' elif s == Status.STATUS_GAME_DRAW: game_end = 'D' else: m = flat(CmdGame.CMD_GAME_MAKE_MOVE, *move, word_size=8) io2.send(m) s = u8(io2.recvn(1)) if s not in (Status.STATUS_SUCCESS, Status.STATUS_GAME_CHECKMATE, Status.STATUS_GAME_DRAW): raise CheckerError('game engine fault: invalid move verdict') if not (io1.recv(1) != s and io1.recvn(6) != m): raise CheckerError('game engine fault: opponent got invalid move') if s == Status.STATUS_GAME_CHECKMATE: game_end = 'W' if order else 'B' elif s == Status.STATUS_GAME_DRAW: game_end = 'D' order = not order if send_msg_after_move == i: self._send_msg(io1 if send_msg_color else io2, msg) valid_pgn += f'\n{{{"W" if send_msg_color else "B"}: "{msg}"}}'.encode() self._is_prev_line_msg = 1 if self._recv_msg(io2 if send_msg_color else io1) != msg.encode(): raise CheckerError('opponent got invalid message') with self._connect() as io: io.send(flat( p8(Cmd.CMD_SRV_VIEW_GAME), room_code, room_pass )) if u8(io.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError("failed to auth in room") pgn = io.recvn(u16(io.recvn(2))) if pgn != valid_pgn: raise CheckerError("invalid PGN") def put(self, flag_id: str, flag: str, vuln: int) -> None: if vuln != 1: raise CheckerInternalError(f"unsupported vuln {vuln}") self._put_vuln(build_vuln_context(flag_id), flag) def get(self, flag_id: str, flag: str, vuln: int) -> None: if vuln != 1: raise CheckerInternalError(f"unsupported vuln {vuln}") self._get_vuln(build_vuln_context(flag_id), flag) def _put_vuln(self, context: VulnContext, flag: str): with self._connect() as io1, self._connect() as io2: self._enter_room(context.room_code, context.room_pass, io1, io2) color = u8(io1.recvn(1)) io2.recvn(1) if not color: io1, io2 = io2, io1 game_end = None order = True for i, move in enumerate(self._pgn_iter(context.pgn_id)): if game_end is not None and not isinstance(move, str): raise CheckerError('game engine fault: invalid game end') if isinstance(move, str): if game_end != move: raise CheckerError('game engine fault: invalid game end') print(json.dumps({'room': context.room_code.decode()})) return if context.comm.after_move == i: self._send_msg(io1 if context.comm.color else io2, f'Such a great move! {flag}') if order: m = flat(CmdGame.CMD_GAME_MAKE_MOVE, *move, word_size=8) io1.send(m) s = u8(io1.recvn(1)) if s not in (Status.STATUS_SUCCESS, Status.STATUS_GAME_CHECKMATE, Status.STATUS_GAME_DRAW): raise CheckerError('game engine fault: invalid move verdict') if not (io2.recv(1) != s and io2.recvn(6) != m): raise CheckerError('game engine fault: opponent got invalid move') if s == Status.STATUS_GAME_CHECKMATE: game_end = 'W' if order else 'B' elif s == Status.STATUS_GAME_DRAW: game_end = 'D' else: m = flat(CmdGame.CMD_GAME_MAKE_MOVE, *move, word_size=8) io2.send(m) s = u8(io2.recvn(1)) if s not in (Status.STATUS_SUCCESS, Status.STATUS_GAME_CHECKMATE, Status.STATUS_GAME_DRAW): raise CheckerError('game engine fault: invalid move verdict') if not (io1.recv(1) != s and io1.recvn(6) != m): raise CheckerError('game engine fault: opponent got invalid move') if s == Status.STATUS_GAME_CHECKMATE: game_end = 'W' if order else 'B' elif s == Status.STATUS_GAME_DRAW: game_end = 'D' order = not order def _get_vuln(self, context: VulnContext, flag: str): with self._connect() as io: io.send(flat( p8(Cmd.CMD_SRV_VIEW_GAME), context.room_code, context.room_pass )) if u8(io.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError("failed to auth in room") pgn_size = u16(io.recvn(2)) pgn = io.recvn(pgn_size) if flag.encode() not in pgn: raise CorruptError("flag not present in PGN") def _connect(self): return arc4tube(self.host, self.port, timeout=TIMEOUT) def _enter_room(self, room_code, room_pass, io1, io2): io1.send(flat(p8(Cmd.CMD_SRV_ENTER_ROOM), room_code)) if u8(io1.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError("failed to create room") sleep(.1) io2.send(flat(p8(Cmd.CMD_SRV_ENTER_ROOM), room_code)) if u8(io2.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError("failed to enter room") io1.send(room_pass) if u8(io1.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError("failed to protect room with password") def _pgn_iter(self, pgn_id: int, pgn_dump: Optional[bytearray] = None): pgn = _load_pgn(pgn_id) res = {"1-0": "W", "0-1": "B", "1/2-1/2": "D"}[pgn.headers["Result"]] board = chess.Board() i = 0 for move in pgn.mainline_moves(): row_from = 7 - (move.from_square // 8) col_from = move.from_square % 8 row_to = 7 - (move.to_square // 8) col_to = move.to_square % 8 m = [row_from, col_from, row_to, col_to] if board.is_kingside_castling(move): m += [MovePayload.MOVE_KING_KINGSIDE_CASTLE, 0] elif board.is_queenside_castling(move): m += [MovePayload.MOVE_KING_QUEENSIDE_CASTLE, 0] elif move.promotion is not None: match move.promotion: case chess.PAWN: pv = Piece.PIECE_PAWN case chess.KNIGHT: pv = Piece.PIECE_KNIGHT case chess.BISHOP: pv = Piece.PIECE_BISHOP case chess.ROOK: pv = Piece.PIECE_ROOK case chess.QUEEN: pv = Piece.PIECE_QUEEN pv |= Color.COLOR_WHITE if board.turn is chess.WHITE else Color.COLOR_BLACK m += [MovePayload.MOVE_PAWN_PROM, pv] else: m += [MovePayload.MOVE_NONE, 0] if pgn_dump is not None: if board.turn is chess.WHITE: i += 1 pgn_dump += f'\n{i}. {board.lan(move)}'.encode() else: if self._is_prev_line_msg: pgn_dump += f'\n{i}... '.encode() pgn_dump += f' {board.lan(move)}'.encode() self._is_prev_line_msg = 0 yield tuple(m) board.push(move) yield res def _send_msg(self, io, msg: str): io.send(flat(p8(CmdGame.CMD_GAME_SEND_MESSAGE), p16(len(msg)), msg.encode())) if u8(io.recvn(1)) != Status.STATUS_SUCCESS: raise CheckerError('failed to send message') def _recv_msg(self, io): io.send(p8(CmdGame.CMD_GAME_RECV_MESSAGE)) n = u8(io.recvn(1)) msg = b'' for i in range(n): msg += io.recvn(u16(io.recvn(2))) return msg def parse_args(): # sigma patch import sys class Namespace: pass n = Namespace() if len(sys.argv) >= 2: n.action = sys.argv[1] if len(sys.argv) >= 3: n.host = sys.argv[2] if len(sys.argv) >= 4: n.flag_id = sys.argv[3] if len(sys.argv) >= 5: n.flag = sys.argv[4] if len(sys.argv) >= 6: n.vuln = int(sys.argv[5]) return n parser = argparse.ArgumentParser(description="Grob 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 = GrobChecker(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: # EOF like errors log(f"DOWN") return EXIT_DOWN return EXIT_OK if __name__ == "__main__": sys.exit(main())