#!/usr/bin/env python3
import sys
from hashlib import sha1
from decimal import Decimal, getcontext
from Crypto.Util import number
from random import randint


def random_bytes(size):
    i = randint(0, 2 ** (8 * size))
    return i.to_bytes(size, sys.byteorder)


def _get_prime(N=512, e=3):
    return number.getStrongPrime(N, e=e)


def cube_root(i):
    getcontext().prec = (i.bit_length() + 7) // 8
    answer = Decimal(i) ** (Decimal(1) / Decimal(3))
    return answer.quantize(Decimal("1."))


def invmod(a, n):
    t = 0
    new_t = 1
    r = n
    new_r = a
    while new_r != 0:
        quotient = r // new_r
        t, new_t = new_t, t - quotient * new_t
        r, new_r = new_r, r - quotient * new_r
    if r > 1:  # a is not invertible
        return None
    if t < 0:
        t += n
    return t


def int_to_bytes(i):
    return i.to_bytes(int_byte_size(i), "big")


def int_byte_size(i):
    return (i.bit_length() + 7) // 8


class rsa:
    def generate_keypair(self, key_size=1024, e=3):
        self.e = e  # Public information provided with both keys
        p = _get_prime(N=key_size // 2, e=e)
        q = _get_prime(N=key_size // 2, e=e)
        self.public_key = p * q

        et = (p - 1) * (q - 1)
        self._private_key = invmod(e, et)

    def import_public_key(self, n, e=3):
        self.e = e
        self.public_key = n

    def encrypt(self, plain_text):
        msg = int.from_bytes(plain_text, "big")
        ct = pow(msg, self.e, self.public_key)
        return ct.to_bytes(int_byte_size(self.public_key), "big")

    def decrypt(self, cipher_text):
        ct = int.from_bytes(cipher_text, "big")
        msg = pow(ct, self._private_key, self.public_key)
        return int_to_bytes(msg)

    def _generate_signature(self, data):
        hasher = sha1()
        mod_size = int_byte_size(self.public_key)
        prefix = b"\x00\x01\xff"
        middle = b"\x00LazyASN.1Struc"
        padding_size = mod_size - len(prefix) - len(middle) - hasher.digest_size
        if padding_size < 0:
            raise ValueError("RSA Key is too small to sign with")
        padding = b"\xff" * padding_size
        hasher.update(data)
        return b"".join([prefix, padding, middle, hasher.digest()])

    def sign(self, data):
        return self.decrypt(self._generate_signature(data))

    def verify(self, data, signature):
        # # Easy and more secure way:
        # return self.encrypt(signature) == self._generate_signature(data)

        # # Hard way:
        decrypted_signature = self.encrypt(signature)

        # Make sure the signature starts with the correct prefix
        prefix = b"\x00\x01\xff"
        if not decrypted_signature.startswith(prefix):
            return False

        # Validate the padding structure badly
        end_of_padding = decrypted_signature.index(b"\x00", len(prefix) - 1)
        expected_padding = b"\xff" * (end_of_padding - len(prefix))
        if decrypted_signature[len(prefix) : end_of_padding] != expected_padding:
            return False

        # Validate the hashing structure used
        middle = b"\x00LazyASN.1Struc"
        end_of_middle = end_of_padding + len(middle)
        if decrypted_signature[end_of_padding:end_of_middle] != middle:
            return False

        # Create the expected hash and extract the found hash
        hash = sha1()
        raw_hash = decrypted_signature[end_of_middle : end_of_middle + hash.digest_size]
        hash.update(data)

        return hash.digest() == raw_hash


def forge_signature(public_key, data):
    key_size = int_byte_size(public_key)
    prefix = b"\x00\x01\xff\x00LazyASN.1Struc"
    hash = sha1(data).digest()
    prefix = prefix + hash

    padding_size = key_size - len(prefix)
    block = prefix + (b"\x00" * padding_size)
    block = int.from_bytes(block, "big")
    if block >= public_key:
        raise ValueError("Needed block would wrap the public key modulus")

    rooted = int(cube_root(block))
    # Confirm when its cubed it will result in the expected prefix
    if not (rooted ** 3).to_bytes(key_size, "big").startswith(prefix):
        raise ValueError("Public key size is too small to forge with")

    return int_to_bytes(rooted)


def test_signing(loops=3):
    for _ in range(loops):
        data = random_bytes(randint(32, 1024))
        private = rsa()
        private.generate_keypair()
        signature = private.sign(data)

        public = rsa()
        public.import_public_key(private.public_key)

        if not public.verify(data, signature):
            return False
    return True


def test_forging():
    data = b"hi mom"  # Data to forge a signature for

    # Create a key pair but only use the public key
    private = rsa()
    private.generate_keypair()

    forged_signature = forge_signature(private.public_key, data)

    public = rsa()
    public.import_public_key(private.public_key)

    if not public.verify(data, forged_signature):
        return False
    return True


if test_signing():
    print("RSA Signing Tests Passed!")
else:
    print("RSA Signing Tests Failed!")

if test_forging():
    print("RSA Forging Tests Passed!")
else:
    print("RSA Forging Tests Failed!")
