#!/usr/bin/sage

from sage.all import *
import struct
import re
from Crypto.PublicKey import RSA
from Crypto.Cipher import AES
from Crypto import Random
import base64

def b64_enc(s):
    return base64.encodebytes(s).decode("ascii")

def b64_dec(s):
    return base64.b64decode(s)

# Our "MPI" format consists of 4-byte integer length l followed by l bytes of binary key
def int_to_mpi(z):
    s = int_to_binary(z)
    return struct.pack('<I',len(s))+s

# Get bytes representation of arbitrary-length long int
def int_to_binary(z):
    z = int(z)
    return z.to_bytes((z.bit_length() + 7) // 8, 'big')

encrypt_header = '-----BEGIN PRETTY BAD ENCRYPTED MESSAGE-----\n'
encrypt_footer = '-----END PRETTY BAD ENCRYPTED MESSAGE-----\n'

# PKCS 7 pad message.
def pad(s,blocksize=AES.block_size):
    n = blocksize-(len(s)%blocksize)
    return s+bytes([n]*n)

# Encrypt plaintext pt using the given rsakey.
# PyCryptodome uses a slightly different API than PyCrypto. Try encrypting the plaintext using the PyCrypto convention, and if this fails, then fall back to the same operation in pure Python.
def rsa_encrypt(rsakey, pt):
    keysize = rsakey.n.bit_length()
    k = ceil(Rational(keysize)/8)-3-32
    EB = '0001' + 'ff'*k + '00' + "%064x"%pt
    m = int(EB, 16)
    try:
        return rsakey.encrypt(m,None)
    except NotImplementedError:
        return pow(m, rsakey.e, rsakey.n)

# Encrypt string s using RSA encryption with AES in CBC mode.
# Generate a 256-bit symmetric key, encrypt it using RSA with PKCS1v1.5 type 1 padding, and prepend the MPI-encoded RSA ciphertext to the AES-encrypted ciphertext of the message.
def encrypt(rsakey,s):
    m = ZZ.random_element(2**256)

    output = int_to_mpi(rsa_encrypt(rsakey, m))

    aeskey = int_to_binary(m)
    iv = Random.new().read(AES.block_size)
    cipher = AES.new(aeskey, AES.MODE_CBC, iv)

    output += iv + cipher.encrypt(pad(s))

    return encrypt_header + b64_enc(output) + encrypt_footer

if __name__=='__main__':
    pubkey = RSA.importKey(open('key.pub').read())
    plaintext = open('hw8.pdf', 'rb').read()

    f = open('hw8.pdf.enc.asc','w')
    f.write(encrypt(pubkey,plaintext))
    f.close()