# PlayCrypt stuff
from playcrypt.new_tools import K_rsa

from playcrypt.games.game_ufcma_sign import GameUFCMASign
from playcrypt.simulator.ufcma_sign_sim import UFCMASignSim

# Bleichenbacher Attack & modulus-reuse stuff

from playcrypt.tools import string_to_int, int_to_string
from Crypto.Cipher import PKCS1_OAEP
from Crypto.Util.number import getPrime
from Crypto.PublicKey import RSA
from Crypto.Hash import SHA
import base64
from math import gcd
import random

def RSAlice_Gen():
    p = getPrime(1024)
    q = getPrime(1024)
    n = p * q
    pk = n
    sk = (p, q)
    return (pk, sk)

def RSAlice_Enc(pk, m):
    e = getPrime(128)
    y = pow(m, e, pk)
    return (e, y)

def RSAlice_Dec(sk, c):
    (p, q) = sk
    (e, y) = c
    n = p * q
    phi_n = (p-1) * (q-1)
    d = pow(e, -1, phi_n)
    m = pow(y, d, n)
    return m

def P1_demo(message):
    (pk, sk) = RSAlice_Gen()
    message_int = string_to_int(message)

    (e1, y1) = RSAlice_Enc(pk, message_int)
    with open("ciphertext1.txt", "w") as f:
        f.write(f"pk = {pk}\n")
        f.write(f"e = {e1}\n")
        f.write(f"y = {y1}")
    
    (e2, y2) = RSAlice_Enc(pk, message_int)
    with open("ciphertext2.txt", "w") as f:
        f.write(f"pk = {pk}\n")
        f.write(f"e = {e2}\n")
        f.write(f"y = {y2}")

def P1():
    token = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA"
    return token

def P1_code():
    pass

def K1():
    (N, p, q, e, d) = K_rsa(k)
    pk = (N, e)
    sk = (N, d)
    return (pk, sk)

def S1(sk, m):
    """
    :param sk: The secret key sk = (N, d) used to sign messages
    :param M: The message to be signed
    :return: return the signature of message M
    """
    (N, d) = sk
    if gcd(N, m) != 1:
        return False
    return pow(m, d, N)

def V1(pk, m, s):
    """
    :param pk: The public key pk = (N, e) used to verify signatures
    :param M: The message to be verified
    :param S: The signature to be verified
    :return: return 1 if the signature is valid, 0 otherwise
    """
    (N, e) = pk
    if gcd(N, m) != 1:
        return 0
    return (1 if (pow(s, e, N) == m) else 0)

"""
Problem 2 (30 points): Present an O(1)-time adversary A1 making ZERO queries to its Sign oracle and achieving Adv^{uf-cma}_DS(A1) = 1.
"""
def A1(Sign, pk):
    pass

# this is Alice's public key, which your forged signatures will be verified against.

publickey = """-----BEGIN PUBLIC KEY-----
MIIBIDANBgkqhkiG9w0BAQEFAAOCAQ0AMIIBCAKCAQEAqtMbmgeGOL5l+sylkG5C
AgAmCmCXmN/KNFuIJaF1cxKMoiZzlqew3DVNF+Xs5rkFkzrflw2MVLY8SQl/qyRO
yHNy68OVwXeAbSIyY/8reUh2AXTm013HVS+LvI6yVOgQ4AwvfbuAPQ4B+nYbkK9G
wgHczJrChPMOaZz7yMBBwwEeonqdeNkuAyAsM/E7UmxCsR3FdMF3vuARLY/+7UJx
wDMFO1LMt5zOrQtK3AKiT4GTv4orBMAZ159ocBawpq6Z5emuI6opGavxLrjTlQgG
KagUNHhQXnQ/+pX6wPuMzWVv21z6L6m3Fbm/bWpyLteftEO7d+vMS8HTDzFQgjN2
bwIBAw==
-----END PUBLIC KEY-----"""


def bleichenbacher(m):
    """
    Problem 3 (60 points): Break the flawed verification function (see `sus_verify` for the source code)!

    This function takes in
    - m: the message you want to forge a signature for, as a pure string.
    
    Your function should return:
    - A forged signature for the message m, which should verify successfully under the RSA key named publickey.
    - THIS FORGED SIGNATURE SHOULD BE ENCODED IN BASE64.
    """

    # ## Python notes:
    # # In Python, if you write x^y, this is NOT exponentiation, it's bitwise XOR.
    # # Instead for exponentiation you would write x**y or pow(x,y)
    # # To compute x^y mod z, you can write x**y % z
    # # Better yet, use three-argument pow: pow(x,y,z)
    # # pow(x,y,z) is much faster than (x**y % z) when y is very large,
    # # because it reduces mod z as it goes, rather than computing x**y first.
    # ##
    
    return None

# You may want to use the integer_nthroot, integer_to_base64, integer_to_bytes, and the bytes_to_integer helper methods below:

# The below code for integer_nthroot, integer_to_base64, integer_to_bytes, bytes_to_integer is taken from SymPy
# Used under fair-use.
# https://raw.githubusercontent.com/sympy/sympy/733da515a7638bba4e08be366bf24c996ad84a61/sympy/core/power.py

from math import sqrt, trunc, log as _log

def integer_nthroot(y, n):
    """
    Return a tuple containing x = floor(y**(1/n))
    and a boolean indicating whether the result is exact (that is,
    whether x**n == y).

    >>> from sympy import integer_nthroot
    >>> integer_nthroot(16,2)
    (4, True)
    >>> integer_nthroot(26,2)
    (5, False)

    """
    y, n = int(y), int(n)
    if y < 0:
        raise ValueError("y must be nonnegative")
    if n < 1:
        raise ValueError("n must be positive")
    if y in (0, 1):
        return y, True
    if n == 1:
        return y, True
    if n == 2:
        x = trunc(sqrt(y))
        rem = y - x * x
        return int(x), not rem
    if n > y:
        return 1, False
    # Get initial estimate for Newton's method. Care must be taken to
    # avoid overflow
    try:
        guess = int(y ** (1.0 / n) + 0.5)
    except OverflowError:
        exp = _log(y, 2) / n
        if exp > 53:
            shift = int(exp - 53)
            guess = int(2.0 ** (exp - shift) + 1) << shift
        else:
            guess = int(2.0**exp)
    if guess > 2**50:
        # Newton iteration
        xprev, x = -1, guess
        while 1:
            t = x ** (n - 1)
            xprev, x = x, ((n - 1) * x + y // t) // n
            if abs(x - xprev) < 2:
                break
    else:
        x = guess
    # Compensate
    t = x**n
    while t < y:
        x += 1
        t = x**n
    while t > y:
        x -= 1
        t = x**n
    return x, t == y



def integer_to_base64(z):
    """Converts an arbitrarily long integer to a big-endian base64 encoding"""
    s = "%x" % z
    s = ("0" * (len(s) % 2)) + s
    s = bytes.fromhex(s)
    return base64.b64encode(s)


def integer_to_bytes(n):
    """Converts an arbitrarily long integer to bytes, big-endian"""
    negative = True if n < 0 else False
    out = list()
    while n > 0:
        out.append(n & 0xFF)
        n >>= 8
    if not negative:
        out.append(0x00)
    return bytes(reversed(out))


def bytes_to_integer(b):
    """Converts big-endian bytes to an arbitrarily long integer"""
    out = 0
    for digit in b:
        out <<= 8
        assert digit >= 0
        out += digit
        assert out >= 0
    return out


## You don't need to (and should not) modify anything below. It's used for testing your code locally, and for helper methods you may want to use.


# This function was adapted from the autograder code of the CSE127 "Cryptography" assignment.
def sus_verify(message, sig):
    m = message.encode("utf-8")
    key = RSA.importKey(publickey)
    sig = bytes_to_integer(base64.b64decode(sig))
    hexd = integer_to_bytes(pow(sig, key.e, key.n)).hex()
    if hexd[:4] != "0001":
        print("Bad beginning")
        return False
    hexd = hexd[4:]
    if hexd[:2] != "ff":
        print("No ff padding bytes found")
        return False
    while hexd.startswith("ff"):
        hexd = hexd[2:]
    asn1_magic = "003021300906052b0e03021a05000414"
    if hexd[: len(asn1_magic)] != asn1_magic:
        print("Bad ASN1 magic")
        return False
    hexd = hexd[len(asn1_magic) :]
    h = SHA.new(message.encode("utf-8")).hexdigest()
    if hexd[: len(h)] != h:
        print("Bad SHA")
        return False
    print("Forged signature passed verification!")
    return True

# This function was adapted from the autograder code of the CSE127 "Cryptography" assignment.
def random_pw(length):
    return "".join(
        random.choices(
            ("abcdefghijklmnopqrstuvwxyz1234567890!~*#%&^$@_-+='?/"), k=length
        )
    )


def test_bleichenbacher():
    print("Checking Bleichenbacher attack...")
    for i in range(1, 15):
        print(f"Testing for len(m) = {i * 5}")
        test_m = random_pw(i * 5)
        forged_signature = bleichenbacher(test_m)
        if not sus_verify(test_m, forged_signature):
            break
    else:
        print("Passed all cases!")

def test_P2():
    g1 = GameUFCMASign(S1, V1, 0, K1)
    s1 = UFCMASignSim(g1, A1)
    print("The advantage of your adversary A1 is approx. " + str(s1.compute_advantage()))

if __name__ == "__main__":
    k = 64
    test_P2()
    test_bleichenbacher()
    print()