from z3 import *

ALPHABET = "SDHACENOIFKXQLBMPJTZURWVGY"

KEY1_INDEXES = [ 3,  4, 28, 29,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15,  5, 17, 18, 19, 20, 21, 22]
KEY2_INDEXES = [ 3,  4,  1,  2, 29, 30, 15, 16, 31, 32, 33, 34, 35, 36, 28, 17, 18, 19, 20, 21, 22]

def encodeString(s):
    result = s[0]
    lastLetterIdx = ALPHABET.index(s[0])

    for i in range(len(s) - 1):
        thisLetterIdx = ALPHABET.index(s[i + 1])
        cypherNum = thisLetterIdx - lastLetterIdx - i
        cypherNum += 284
        modNum = cypherNum % 26
        encChar = chr(modNum + 48)

        result += encChar

        lastLetterIdx = thisLetterIdx

    return result

def encodedStringToNumber(s):
    number = 0
    letterMulti = 1
    for i in range(len(s), 0, -1):
        if(i < len(s)):
            letterMulti = letterMulti * 19

        letter = ord(s[i-1])
        letter -= 48
        letter *= letterMulti

        number += letter

    return number

def is_checksum_valid(s):
    sum_val = sum([ord(c) for c in s[1:]])
    mod_val = sum_val % 25
    return ALPHABET[mod_val] == s[0]

def pad(val, pad_num):
    return str(val).rjust(pad_num, '0')

def get_key1(s):
    res = encodeString(s)
    if not is_checksum_valid(res):
        raise Exception("Invalid checksum for key {}".format(s))
    p1 = pad(encodedStringToNumber(res[1:13]), 15)
    p2 = pad(encodedStringToNumber(res[13:]), 15)
    return p1 + p2

def get_key2(s):
    res = encodeString(s)
    if not is_checksum_valid(res):
        raise Exception("Invalid checksum for key {}".format(s))
    p1 = pad(encodedStringToNumber(res[1:12]), 14)
    p2 = pad(encodedStringToNumber(res[12:23]), 14)
    p3 = pad(encodedStringToNumber(res[23:]), 7)
    return p1 + p2 + p3

def extract_indexes_from_key(k, indexes):
    res = ""
    for index in indexes:
        res += k[index - 1]
    return res

def is_valid(key1, key2):
    a = extract_indexes_from_key(key1, KEY1_INDEXES)
    b = extract_indexes_from_key(key2, KEY2_INDEXES)
    if a != b:
        return False
    if key2[5:8] != "999":
        return False
    return True

if __name__ == "__main__":
    initial1 = "ALVMGRCTQQEYMNAHDANKPGKPO"
    res = encodeString(initial1)
    assert(res == "A87>4A26;@8833898:::;<?8B")
    assert(is_checksum_valid(res) == True)
    p1 = pad(encodedStringToNumber(res[1:13]), 15)
    assert(p1 == "979440403325666")
    p2 = pad(encodedStringToNumber(res[13:]), 15)
    assert(p2 == "401568800000018")
    key1 = p1 + p2
    assert(key1 == "979440403325666401568800000018")
    assert(key1 == get_key1(initial1))
    assert(extract_indexes_from_key(key1, KEY1_INDEXES) == "940104033256664015688")

    initial2 = "VABCDEFGHIJKLMNOPQRSTUVWXYZABK"
    res = encodeString(initial2)
    assert(res == "V48<BHG7EFH7@>2;B4@;G661@>C88B")
    assert(is_checksum_valid(res) == True)
    p1 = pad(encodedStringToNumber(res[1:12]), 14)
    assert(p1 == "27326846177342")
    p2 = pad(encodedStringToNumber(res[12:23]), 14)
    assert(p2 == "102659346508195")
    p3 = pad(encodedStringToNumber(res[23:]), 7)
    assert(p3 == "88621338")
    key2 = p1 + p2 + p3
    assert(key2 == "2732684617734210265934650819588621338")
    assert(key2 == get_key2(initial2))
    assert(extract_indexes_from_key(key2, KEY2_INDEXES) == "322758108621339265934")


    INPUT_LEN = 30
    user_input = [BitVec("{}".format(i), 8) for i in range(INPUT_LEN)]
    solver = Solver()
    for i in range(INPUT_LEN):
        solver.add(user_input[i] >= ord('0'))
        solver.add(user_input[i] <= ord('~'))
    solver.add(is_valid(get_key1(initial1), get_key2(user_input)) == True)

    if solver.check() == sat:
        res = ""
        m = solver.model()
        print(m)