130 lines
3.7 KiB
Python
130 lines
3.7 KiB
Python
# built-in dependencies
|
|
import math
|
|
import random
|
|
from typing import Optional
|
|
|
|
# 3rd party dependencies
|
|
import sympy
|
|
|
|
# project dependencies
|
|
from lightdsa.interfaces.signatures import Signature
|
|
from lightdsa.commons.transformation import hashify
|
|
from lightdsa.commons.logger import Logger
|
|
|
|
logger = Logger(module="lightdsa/algorithms/rsa.py")
|
|
|
|
DEFAULT_KEY_SIZE = 3072
|
|
|
|
|
|
class RSA(Signature):
|
|
def __init__(
|
|
self,
|
|
keys: Optional[dict] = None,
|
|
key_size: Optional[int] = None,
|
|
):
|
|
"""
|
|
RSA
|
|
[1] https://sefiks.com/2018/05/21/the-math-behind-rsa-algorithm/
|
|
[2] https://sefiks.com/2023/03/06/a-step-by-step-partially-homomorphic-encryption-example-with-rsa-in-python/
|
|
"""
|
|
self.key_size = key_size
|
|
self.keys = keys or self.generate_keys(key_size or DEFAULT_KEY_SIZE)
|
|
|
|
if key_size is None:
|
|
self.hash_algorithm = "sha256"
|
|
elif key_size <= 1024:
|
|
self.hash_algorithm = "sha1"
|
|
elif key_size <= 2048:
|
|
self.hash_algorithm = "sha224"
|
|
elif key_size <= 3072:
|
|
self.hash_algorithm = "sha256"
|
|
elif key_size <= 7680:
|
|
self.hash_algorithm = "sha384"
|
|
else:
|
|
self.hash_algorithm = "sha512"
|
|
|
|
def generate_keys(self, key_size: int) -> dict:
|
|
"""
|
|
Generate public and private keys of RSA cryptosystem
|
|
Args:
|
|
key_size (int): key size in bits
|
|
Returns:
|
|
keys (dict): having private_key and public_key keys
|
|
e.g. keys = {
|
|
"private_key": {
|
|
"d": int
|
|
},
|
|
"public_key": {
|
|
"n": int,
|
|
"e": int
|
|
}
|
|
}
|
|
"""
|
|
keys = {}
|
|
keys["private_key"] = {}
|
|
keys["public_key"] = {}
|
|
|
|
while True:
|
|
try:
|
|
# picking a prime modulus p and q
|
|
p = sympy.randprime(200, 2 ** int(key_size / 2) - 1)
|
|
q = sympy.randprime(200, 2 ** int(key_size / 2) - 1)
|
|
|
|
assert isinstance(p, int)
|
|
assert isinstance(q, int)
|
|
|
|
n = p * q
|
|
phi = (p - 1) * (q - 1)
|
|
|
|
# select public exponent e
|
|
while True:
|
|
e = random.randint(1, phi - 1)
|
|
if math.gcd(e, n) == 1:
|
|
break
|
|
|
|
d = pow(e, -1, phi)
|
|
break
|
|
except:
|
|
pass
|
|
|
|
keys["public_key"]["n"] = n
|
|
keys["public_key"]["e"] = e
|
|
keys["private_key"]["d"] = d
|
|
return keys
|
|
|
|
def sign(self, message: int) -> int:
|
|
"""
|
|
Sign a message with RSA
|
|
Args:
|
|
message (int): message to sign
|
|
Returns:
|
|
signature (int): signature
|
|
"""
|
|
n = self.keys["public_key"]["n"]
|
|
|
|
hashed_message = hashify(message, algorithm=self.hash_algorithm)
|
|
|
|
d = self.keys["private_key"]["d"]
|
|
c = pow(hashed_message, d, n)
|
|
|
|
return c
|
|
|
|
def verify(self, message: int, signature: int) -> bool:
|
|
"""
|
|
Verify a message with RSA
|
|
Args:
|
|
message (int): message to verify
|
|
signature (int): signature
|
|
Returns:
|
|
bool: True if signature is valid, False otherwise
|
|
"""
|
|
n = self.keys["public_key"]["n"]
|
|
e = self.keys["public_key"]["e"]
|
|
hashed_message_prime = pow(signature, e, n)
|
|
hashed_message = hashify(message, algorithm=self.hash_algorithm)
|
|
|
|
if hashed_message_prime != hashed_message:
|
|
raise ValueError("Signature is invalid")
|
|
|
|
return True
|