Never Ending Crypto Redux solver - from a TUCTF 2018 challenge

solve.py

from pwn import *
from itertools import cycle
import string
import codecs
import traceback

MORSE = {
    '.-': 'A',
    '-...': 'B',
    '-.-.': 'C',
    '-..': 'D',
    '.': 'E',
    '..-.': 'F',
    '--.': 'G',
    '....': 'H',
    '..': 'I',
    '.---': 'J',
    '-.-': 'K',
    '.-..': 'L',
    '--': 'M',
    '-.': 'N',
    '---': 'O',
    '.--.': 'P',
    '--.-': 'Q',
    '.-.': 'R',
    '...': 'S',
    '-': 'T',
    '..-': 'U',
    '...-': 'V',
    '.--': 'W',
    '-..-': 'X',
    '-.--': 'Y',
    '--..': 'Z',
}

def decrypt(level, message):
    if level in [2, 5, 7]:
        message = ''.join('.-'['-.'.index(x)] if x in '.-' else x for x in message)
    message = ''.join(map(MORSE.get, message.split()))
    if level in [1, 4, 5]:
        message = codecs.decode(message, 'rot13')
    if level in [3, 4, 5, 7]:
        indices = sum((range(x, len(message), 3) for x in range(3)), [])
        decoded = [None] * len(message)
        for i, c in zip(indices, message):
            decoded[i] = c
        message = ''.join(decoded)
    if level in [6, 7]:
        decoded = ''
        for c, k in zip(message, cycle('TUCTF' if level == 6 else 'HAVEFUN')):
            C, K = map(string.ascii_uppercase.index, (c, k))
            M = (C - K) % 26
            decoded += string.ascii_uppercase[M]
        message = decoded
    return message

s = remote('18.223.156.26', 12345)

level = 0
while level < 7:
    test_vector = 'A' * 50
    s.recvuntil('Level ')
    level = int(s.recvline())
    s.recvuntil('Give me text:')
    s.sendline(test_vector)
    s.recvuntil(' encrypted is ')
    encrypted = s.recvline().strip()
    try:
        decrypted = decrypt(level, encrypted)
    except Exception as e:
        print(e)
        decrypted = None

    if decrypted != test_vector:
        print encrypted
        print decrypted
        print test_vector
        break

    s.recvline()
    p = log.progress('Answering level ' + str(level))
    words = 0
    while True:
        prompt = s.recvline()
        if 'Decrypt ' not in prompt:
            if prompt != 'Congratulations! You beat Level ' + str(level) + '!\n':
                print prompt
            break
        question = prompt.strip()[8:]
        answer = decrypt(level, question)
        s.sendline(answer)
        s.recvline()
        grade = s.recvline().strip()
        s.recvline()
        words += 1
        p.status(str(words) + ' words: ' + grade)
    p.success()
s.interactive()