# generate.py
# Apply custom ROT encryption using QWERTY alphabet

def rot_custom(text: str, shift: int, alphabet: str = "qwertyuiopasdfghjklzxcvbnm") -> str:
    alphabet_lower = alphabet.lower()
    alphabet_upper = alphabet.upper()
    result = []

    for char in text:
        if char in alphabet_lower:
            idx = (alphabet_lower.index(char) + shift) % len(alphabet_lower)
            result.append(alphabet_lower[idx])
        elif char in alphabet_upper:
            idx = (alphabet_upper.index(char) + shift) % len(alphabet_upper)
            result.append(alphabet_upper[idx])
        else:
            result.append(char)

    return "".join(result)


if __name__ == "__main__":
    custom_keyboard = "qwertyuiopasdfghjklzxcvbnm"
    flag = "CTF{custom_rot_with_keyboard_layout}"
    shift =  8

    encrypted = rot_custom(flag, shift, custom_keyboard)
    print(encrypted)