SECCON 2022 CTF Quals: noiseccon

noiseccon.py

import base64
import io

import numpy as np
from PIL import Image
from pwn import *

flag_length = 21
pad_length = 8
total_length = pad_length * 2 + flag_length

height = width = 256
n_img = 20
n_decode = 16

def get_img(idx):
    ##### set scale #####
    scaleX = 1 << (8 * (total_length - pad_length - idx) - 4)
    scaleY = 1 << (8 * (total_length - pad_length - idx))

    ##### remote #####
    rmt = remote('noiseccon.seccon.games', 1337)
    rmt.recv()
    rmt.recv()
    rmt.sendline(str(scaleX).encode())
    rmt.recv()
    rmt.sendline(str(scaleY).encode())
    img_b64 = rmt.recvline().decode()

    ##### convert to PIL.Image #####
    img_raw = base64.b64decode(img_b64)
    img = Image.open(io.BytesIO(img_raw))
    return img

def decode(img, verbose=False):
    ##### define float-th index of np.array #####
    def at(arr, i_float):
        q, r = int(i_float), i_float - int(i_float)
        if r == 0: return arr[q]
        else: return (1 - r) * arr[q] + r * arr[q + 1]
    
    ##### calculate RSS #####
    pix = np.array(img)[:, :, 0]
    dist = np.zeros((n_decode, n_decode))
    for i in range(height // n_img * n_decode):
        for j in range(width // n_img * n_decode):
            dist[-i % n_decode][-j % n_decode] += (at(at(pix,
                i * n_img / n_decode),
                j * n_img / n_decode) - 128) ** 2
    return chr(np.argmin(dist))

flag = ''.join(decode(get_img(i)) for i in range(flag_length))
print(flag)