def test_run_check_exception_timeout(self, stdout_io):
# pylint: disable=abstract-method
class ExceptionChecker(checkerlib.BaseChecker):
def place_flag(self, tick):
raise socket.timeout()
checkerlib.run_check(ExceptionChecker)
self.assertEqual(stdout_io.getvalue(), 'Check result: DOWN\n')
def test_run_check_exception(self):
# pylint: disable=abstract-method
class ExceptionChecker(checkerlib.BaseChecker):
def place_flag(self, tick):
raise ValueError()
with self.assertRaises(ValueError):
checkerlib.run_check(ExceptionChecker)
def test_run_check_tick0(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.OK
MockChecker.flag_mock.return_value = CheckResult.OK
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(0)
MockChecker.service_mock.assert_called_once()
MockChecker.flag_mock.assert_called_once_with(0)
self.assertEqual(stdout_io.getvalue(), 'Check result: OK\n')
def test_run_check_flag_fail(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.OK
MockChecker.flag_mock.return_value = CheckResult.FLAG_NOT_FOUND
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(10)
MockChecker.service_mock.assert_called_once()
MockChecker.flag_mock.assert_called_once_with(10)
self.assertEqual(stdout_io.getvalue(), 'Check result: FLAG_NOT_FOUND\n')
def test_run_check_service_fail(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.FAULTY
MockChecker.flag_mock.return_value = CheckResult.OK
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(10)
MockChecker.service_mock.assert_called_once()
MockChecker.flag_mock.assert_not_called()
self.assertEqual(stdout_io.getvalue(), 'Check result: FAULTY\n')
def test_run_check_tick3(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.OK
MockChecker.flag_mock.return_value = CheckResult.OK
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(3)
MockChecker.service_mock.assert_called_once()
self.assertEqual(MockChecker.flag_mock.call_count, 4)
self.assertEqual(MockChecker.flag_mock.call_args_list, [call(3), call(2), call(1), call(0)])
self.assertEqual(stdout_io.getvalue(), 'Check result: OK\n')
def test_run_check_not_recovering(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.OK
MockChecker.flag_mock.side_effect = [CheckResult.OK, CheckResult.OK, CheckResult.FLAG_NOT_FOUND,
CheckResult.OK, CheckResult.FAULTY]
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(10)
MockChecker.service_mock.assert_called_once()
self.assertEqual(MockChecker.flag_mock.call_count, 5)
self.assertEqual(stdout_io.getvalue(), 'Check result: FAULTY\n')
def test_run_check_basic(self, stdout_io):
MockChecker.reset_mocks()
MockChecker.place_mock.return_value = CheckResult.OK
MockChecker.service_mock.return_value = CheckResult.OK
MockChecker.flag_mock.return_value = CheckResult.OK
checkerlib.run_check(MockChecker)
MockChecker.place_mock.assert_called_once_with(10)
MockChecker.service_mock.assert_called_once()
self.assertEqual(MockChecker.flag_mock.call_count, 6)
self.assertEqual(MockChecker.flag_mock.call_args_list, [call(10), call(9), call(8), call(7), call(6),
call(5)])
self.assertEqual(stdout_io.getvalue(), 'Check result: OK\n')
#!/usr/bin/env python3
from ctf_gameserver import checkerlib
class TestChecker(checkerlib.BaseChecker):
def place_flag(self, tick):
raise Exception('This is fine')
def check_service(self):
return checkerlib.CheckResult.OK
def check_flag(self, tick):
return checkerlib.CheckResult.OK
if __name__ == '__main__':
checkerlib.run_check(TestChecker)
pwn.log.info(
f"Overall duration for check_service: {int(time.time() - start_time)}s"
)
return checkerlib.CheckResult.OK
def check_flag(self, tick: int) -> checkerlib.CheckResult:
"""Tries to retrieve a flag."""
start_time = time.time()
interaction = Interaction(self.ip)
flag = checkerlib.get_flag(tick)
key = self._get_key_for_flag(flag)
train = interaction.get_train(key)
pwn.log.info(
f"Overall duration for check_flag: {int(time.time() - start_time)}s"
)
if not train:
return checkerlib.CheckResult.FLAG_NOT_FOUND
if len(train.wagons) != 1:
return checkerlib.CheckResult.FLAG_NOT_FOUND
if train.wagons[0].content != flag:
return checkerlib.CheckResult.FLAG_NOT_FOUND
return checkerlib.CheckResult.OK
if __name__ == "__main__":
checkerlib.run_check(MarsexpressChecker)
logging.warning('Received non-UTF-8 data: %s', repr(resp))
return checkerlib.CheckResult.FAULTY
if resp != flag:
logging.warning('Received wrong response to GET command')
return checkerlib.CheckResult.FLAG_NOT_FOUND
conn.close()
return checkerlib.CheckResult.OK
def connect(ip):
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.connect((ip, 9999))
return sock
def recv_line(conn):
received = b''
while not received.endswith(b'\n'):
new = conn.recv(1024)
if len(new) == 0:
if not received.endswith(b'\n'):
raise EOFError('Unexpected EOF')
break
received += new
return received.decode('utf-8').rstrip()
if __name__ == '__main__':
checkerlib.run_check(ExampleChecker)
