118 lines
3.6 KiB
Python
118 lines
3.6 KiB
Python
#===========================================================================================
|
|
# _____ _____ _ _ _ __ __ ___ ____ _____ _____ ____ _ _
|
|
# |_ _| ____| \ | | / \ | \/ |/ _ \| _ \_ _| ____/ ___| | | |
|
|
# | | | _| | \| | / _ \ | |\/| | | | | |_) || | | _|| | | |_| |
|
|
# | | | |___| |\ |/ ___ \| | | | |_| | _ < | | | |__| |___| _ |
|
|
# |_| |_____|_| \_/_/ \_\_| |_|\___/|_| \_\|_| |_____\____|_| |_|
|
|
#
|
|
# RSA Utils V1 (Python 3.6)
|
|
#
|
|
# Sources:
|
|
# https://nitratine.net/blog/post/asymmetric-encryption-and-decryption-in-python/
|
|
# https://stackoverflow.com/questions/8933237/how-to-find-if-directory-exists-in-python
|
|
# https://stackoverflow.com/questions/273192/how-can-i-safely-create-a-nested-directory
|
|
#
|
|
#===========================================================================================
|
|
|
|
# Basic cryptography tools
|
|
from cryptography.hazmat.backends import default_backend
|
|
from cryptography.hazmat.primitives.asymmetric import rsa
|
|
|
|
# To format created key obj into text to be used/exported
|
|
from cryptography.hazmat.primitives import serialization
|
|
|
|
# To get user home directory
|
|
from pathlib import Path
|
|
|
|
# Encryption
|
|
from cryptography.hazmat.primitives import hashes
|
|
from cryptography.hazmat.primitives.asymmetric import padding
|
|
|
|
def generate_private_key(host):
|
|
# Generate Private key
|
|
private_key = rsa.generate_private_key(
|
|
public_exponent=65537,
|
|
key_size=2048,
|
|
backend=default_backend()
|
|
)
|
|
save_key(private_key, 0, host)
|
|
return private_key
|
|
|
|
|
|
def generate_public_key(host, private_key):
|
|
# Get Public key from Private
|
|
public_key = private_key.public_key()
|
|
save_key(public_key, 1, host)
|
|
return public_key
|
|
|
|
|
|
def private_key_serializer(private_key):
|
|
# Key serialization
|
|
private_key_pem = private_key.private_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PrivateFormat.PKCS8,
|
|
encryption_algorithm=serialization.NoEncryption()
|
|
)
|
|
return private_key_pem
|
|
|
|
|
|
def public_key_serializer(public_key):
|
|
# Key serialization
|
|
public_key_pem = public_key.public_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo
|
|
)
|
|
#print(public_key_pem)
|
|
return public_key_pem
|
|
|
|
|
|
def save_key(key, key_type, host):
|
|
root_file_path = '.auth_server_test/'
|
|
if host == 0:
|
|
path = root_file_path + 'server'
|
|
elif host == 1:
|
|
path = root_file_path + 'client'
|
|
keys_dir_path = Path.home() / path
|
|
# Path exists ? If not we will create it
|
|
keys_dir_path.mkdir(exist_ok=True)
|
|
if(key_type == 0):
|
|
key_pem = private_key_serializer(key)
|
|
key_path = keys_dir_path / 'id_rsa'
|
|
elif(key_type == 1):
|
|
key_pem = public_key_serializer(key)
|
|
key_path = keys_dir_path / 'id_rsa.pub'
|
|
else:
|
|
print("ERROR: Key type error")
|
|
exit()
|
|
|
|
with open(str(key_path), 'wb') as f:
|
|
f.write(key_pem)
|
|
|
|
|
|
def generate_keys(host):
|
|
private_key = generate_private_key(host)
|
|
public_key = generate_public_key(host,private_key)
|
|
return (private_key,public_key)
|
|
|
|
def encrypt_msg(msg,public_key):
|
|
enc_msg = public_key.encrypt(
|
|
msg,
|
|
padding.OAEP(
|
|
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
|
algorithm=hashes.SHA256(),
|
|
label=None
|
|
)
|
|
)
|
|
return enc_msg
|
|
|
|
def decrypt_msg(enc_msg,private_key):
|
|
msg = private_key.decrypt(
|
|
enc_msg,
|
|
padding.OAEP(
|
|
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
|
algorithm=hashes.SHA256(),
|
|
label=None
|
|
)
|
|
)
|
|
return msg
|