Generate passcodes for Small Soldiers (PlayStation)

small_soldiers.py

SYMBOLS = ['X', 'Square', 'Circle', 'Triangle']
SYMBOL_MAP = {s: i for i, s in enumerate(SYMBOLS)}

DIFFICULTIES = {0: "Easy", 1: "Normal", 2: "Hard"}


def decode_passcode(passcode):
    if len(passcode) != 8 or any(s not in SYMBOL_MAP for s in passcode):
        return False, None

    for i, symbol in enumerate(passcode):
        val |= (SYMBOL_MAP[symbol] << (i * 2))

    level = val & 0xF
    ammo = (val & 0xF0) >> 4
    difficulty = (val & 0x300) >> 8
    checksum = (val & 0xFC00) >> 10

    if not (
        level < 0xE and 
        difficulty < 3 and 
        ammo < 8 and 
        (level + ammo + difficulty) == checksum
    ):
        return False, None

    return True, {
        "level": level,
        "difficulty": DIFFICULTIES.get(difficulty),
        "ammo_bits": bin(ammo)[2:].zfill(3),
        "green_ammo": (ammo & 1) != 0,
        "blue_ammo": (ammo & 2) != 0,
        "red_ammo": (ammo & 4) != 0
    }

def encode_passcode(level, difficulty, ammo):
    checksum = level + difficulty + ammo
    
    val = (level & 0xF)
    val |= (ammo & 0x7) << 4
    val |= (difficulty & 0x3) << 8
    val |= (checksum & 0x3F) << 10

    passcode = []
    for i in range(8):
        symbol_idx = (val >> (i * 2)) & 0x3
        passcode.append(SYMBOLS[symbol_idx])
    
    return ', '.join(passcode)


if __name__ == '__main__':
    print(
        'Difficulty',
        'Level',
        'Passcode',
        sep=';'
    )
    for difficulty_index, difficulty_name in DIFFICULTIES.items():
        for level_index in range(0xe):
            passcode = encode_passcode(level_index, difficulty_index, 0b111)
            is_valid, passcode_data = decode_passcode(passcode)
            print(
                difficulty_name,
                level_index + 1,
                passcode,
                sep=';'
            )