import sys
from math import log, ceil, floor
from fractions import Fraction


e = 3
d = 49245850836238243386848117224834103046337172957950760944544575720018018155267
n = 73868776254357365080272175837251154570062235041494422684546086388903725268033


def bytes_needed(n: int) -> int:
    if n == 0:
        return 1
    return int(log(n, 256)) + 1


def add_padding(m: int, k: int) -> int:
    padding_size = 11
    from_len = bytes_needed(m)
    assert(from_len + padding_size <= k)
    padding_str_len = k - 3 - from_len
    r = [0x0, 0x2]
    for _ in range(2, padding_str_len + 2):
        r.append(0xff)
    r.append(0x0)
    r += m.to_bytes(from_len, 'big')
    assert(len(r) == k)
    r = int.from_bytes(r, byteorder='big')
    return r


def remove_padding(m: int, k: int) -> int:
    v = m.to_bytes(k, 'big')
    assert(v[0] == 0x0 and v[1] == 0x2)
    i = 2
    while v[i] == 0xff:
        i = i + 1
    assert(v[i] == 0x0)
    return int.from_bytes(v[i + 1:], 'big')


def oracle(c: int) -> bool:
    m = decrypt(c)
    v = m.to_bytes(bytes_needed(n), 'big')
    return v[0] == 0x0 and v[1] == 0x2


def encrypt(m: int) -> int:
    return pow(m, e, n)


def decrypt(c: int) -> int:
    return pow(c, d, n)


def test():
    m = int.from_bytes(b"kick it, CC", byteorder='big')
    assert(m == 129852745126415640677073731)

    k = bytes_needed(n)
    assert(k == 32)

    padded_m = 5300541194335152988749892502228755547482451611528547105226896651010982723
    assert(padded_m == add_padding(m, k))

    c_new = encrypt(padded_m)
    c = 71554147358804792877798821486588152314859921438911236615156507964101619628630
    assert(c == c_new)

    assert(m == remove_padding(add_padding(m, k), k))
    print('[tests passed]')


def main():
    test()

    k = bytes_needed(n)
    m = int.from_bytes(b"kick it, CC", byteorder='big')
    m_orig = m
    c = encrypt(add_padding(m, k))
    assert(m == remove_padding(decrypt(c), k))

    B = pow(2, 8 * (k - 2))
    # Step 1: Blinding (not necessary because already PKCS conforming).
    i = 1
    s = 1
    m = [(2 * B, 3 * B - 1)]
    assert(oracle(c))

    # Step 2: Searching for PKCS conforming messages.
    while True:
        # print("========")
        # print(f"{i=} {s=} {m=}")
        if i == 1:
            # Step 2.a: Starting the search.
            s = ceil(n / (3 * B))
            while not oracle(c * pow(s, e, n) % n):
                s += 1
        elif len(m) > 1:
            # Step 2.b: Searching with more than one interval left.
            raise Exception("Not implemented")
        elif len(m) == 1:
            # Step 2.c: Searching with one interval left.
            a, b = m[0]
            found = False
            r = ceil(Fraction(2 * (b * s - 2 * B), n))
            while not found:
                s_lower = Fraction(2 * B + r * n, b)
                s_upper = Fraction(3 * B + r * n, a)
                s_lower_ceil = ceil(s_lower)
                s_upper_ceil = ceil(s_upper)
                for s in range(s_lower_ceil, s_upper_ceil):
                    if oracle(c * pow(s, e, n) % n):
                        found = True
                        break
                r += 1
        else:
            raise Exception("No intervals?")

        # Step 3: Narrowing the set of solutions.
        m_new = []
        for (a, b) in m:
            lower = Fraction((a * s - 3 * B + 1), n)
            upper = Fraction((b * s - 2 * B), n)
            lower_ceil = ceil(lower)
            upper_ceil = ceil(upper)

            # If upper is an integer we have to increment it to include it into
            # the range.
            if upper == upper_ceil:
                upper_ceil += 1

            for r in range(lower_ceil, upper_ceil):
                a_new = max(a, ceil(Fraction(2 * B + r * n, s)))
                b_new = min(b, floor(Fraction(3 * B - 1 + r * n, s)))
                m_new.append((a_new, b_new))
        m = m_new

        # Step 4: Computing the solutions.
        if len(m) == 1 and m[0][0] == m[0][1]:
            m = m[0][0]
            break
        i = i + 1

    assert(m == 5300541194335152988749892502228755547482451611528547105226896651010982723)
    m = remove_padding(m, k)
    assert(m == m_orig)
    print("[okay] Challenge 47: Bleichenbacher's PKCS 1.5 Padding Oracle (Simple Case)")


main()