467 lines
15 KiB
Python
Executable File
467 lines
15 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
|
|
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())
|