from pwn import *

def get_word_lengths(raw_bytes):
    # 1. Decode the bytes to a string
    try:
        decoded_output = raw_bytes.decode('utf-8')
    except UnicodeDecodeError:
        # Fallback if bytes are messy, though the provided input is valid utf-8
        decoded_output = raw_bytes.decode('utf-8', errors='ignore')

    print("--- Decoded Output ---")
    print(decoded_output)
    print("----------------------\n")

    # 2. Find the line containing "Structure:"
    structure_line = None
    for line in decoded_output.splitlines():
        if "Structure:" in line:
            structure_line = line
            break
    
    if not structure_line:
        return "No 'Structure' line found."

    # 3. Extract the pattern part (remove "Structure: " prefix)
    # The split ensures we get everything after the colon and space
    pattern_str = structure_line.split("Structure:")[1].strip()

    # 4. Split by the underscore '_' which acts as the delimiter between words
    word_segments = pattern_str.split('_')

    # 5. Calculate lengths of each segment
    word_lengths = [len(segment) for segment in word_segments]
    
    return word_lengths

with remote("chall.polygl0ts.ch", 6052) as p:
    structure = p.read()
    lens = get_word_lengths(structure)

    with open("./word_list.txt", "r") as f:
        words = f.readlines()

    word_dict = {}

    for word in words:
        if word_dict[len(word)]:

            word_dict[len(word)].append(word)
        else:
            word_dict[len(word)] = [word]

    for len in lens:
        print(f"Num of words with len {len} in wordlist: {len(word_dict[len])}")