utils.py

import base64
import hashlib

from Cryptodome import Random
from Cryptodome.Cipher import PKCS1_v1_5
from Cryptodome.Math.Numbers import Integer
from Cryptodome.Util.number import ceil_div, bytes_to_long, long_to_bytes, size
from Cryptodome.Util.py3compat import bchr, bord, b, bstr

'''
pycrypto doesn't allow non hash objects in Cryptodome.Signature.PKCS115_SigScheme.
This modified PKCS115_Cipher function encrypts messages using (d,n) instead of (e,n)
and using 0xFF when padding rather than random bytes.
It's equivalent to the RSA/ECB/PKCS1Padding when encrypting data using private keys in Java.
'''


def _encrypt(key, message):
    mod_bits = size(key.n)
    k = ceil_div(mod_bits, 8)
    m_len = len(message)

    # Step 1
    if m_len > k - 11:
        raise ValueError("Plaintext is too long.")
    # Step 2a
    ps = []
    while len(ps) != k - m_len - 3:
        new_byte = bchr(0xFF)
        if bord(new_byte[0]) == 0x00:
            continue
        ps.append(new_byte)
    ps = b("").join(ps)
    assert (len(ps) == k - m_len - 3)
    # Step 2b
    em = b('\x00\x01') + ps + bchr(0x00) + message
    # Step 3a (OS2IP)
    em_int = bytes_to_long(em)
    # Step 3b (RSAEP)
    if not 0 < em_int < key._n:
        raise ValueError("Plaintext too large")
    m_int = int(pow(Integer(em_int), key._d, key._n))
    # Step 3c (I2OSP)
    c = long_to_bytes(m_int, k)
    return c


def _decrypt(key, ciphertext):
    modBits = size(key.n)
    k = ceil_div(modBits, 8)  # Convert from bits to bytes

    # Step 1
    if len(ciphertext) != k:
        raise ValueError("Ciphertext with incorrect length.")
    # Step 2a (O2SIP)
    ct_int = bytes_to_long(bstr(ciphertext))
    # Step 2b (RSADP)
    m_int = pow(Integer(ct_int), key._e, key._n)
    # Complete step 2c (I2OSP)
    em = long_to_bytes(m_int, k)
    # Step 3
    sep = em.find(bchr(0x00), 2)
    # Step 4
    return em[sep + 1:]


def rsa_encrypt(key, data):
    b = data.encode('utf-8')
    s = b''.join([_encrypt(key, b[i:i + 117]) for i in range(0, len(b), 117)])
    return base64.b64encode(s)


def rsa_encrypt_public(key, data):
    b = data.encode('utf-8')
    cipher = PKCS1_v1_5.new(key)
    s = b''.join([cipher.encrypt(b[i:i + 117]) for i in range(0, len(b), 117)])
    return base64.b64encode(s)


def rsa_decrypt(key, data):
    sentinel = Random.new().read(35)
    cipher = PKCS1_v1_5.new(key)
    b = base64.b64decode(data)
    s = b''.join([cipher.decrypt(b[i:i + 128], sentinel) for i in range(0, len(b), 128)])
    return s.decode('utf-8')


def rsa_decrypt_public(key, data):
    b = base64.b64decode(data)
    s = b''.join([_decrypt(key, b[i:i + 128]) for i in range(0, len(b), 128)])
    return s.decode('utf-8')


def md5_digest(s):
    return hashlib.md5(s.encode('utf-8')).hexdigest()