Merge branch 'main' of https://github.com/uminho-mei-es/2324-G05

TPs/TP04/cfich_nike.py

@@ -0,0 +1,152 @@
+import argparse
+import os
+
+from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.asymmetric import dh
+from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
+from cryptography.hazmat.primitives.kdf.hkdf import HKDF
+import cryptography.hazmat.primitives.serialization as serialization
+from cryptography.hazmat.primitives.serialization import (
+    Encoding,
+    NoEncryption,
+    PrivateFormat,
+    PublicFormat,
+)
+
+def setup(user):
+    parameters = dh.generate_parameters(generator=2, key_size=512)
+    private_key = parameters.generate_private_key()
+    public_key = private_key.public_key()
+
+    skey = open(f"{user}.sk", 'wb')
+    pkey = open(f"{user}.pk", 'wb')
+    skey.write(private_key.private_bytes(encoding=Encoding.PEM, format=PrivateFormat.PKCS8, encryption_algorithm=NoEncryption()))
+    pkey.write(public_key.public_bytes(encoding=Encoding.PEM, format=PublicFormat.SubjectPublicKeyInfo))
+    skey.close()
+    pkey.close()
+
+
+def encrypt(user,filename):
+    peer_pkey_file = open(f"{user}.pk", 'rb')
+    if peer_pkey_file is None:
+        print(f"No public key found for {user}")
+        return
+
+    file_r = open(filename, 'rb')
+    if file_r is None:
+        print(f"File {filename} not found")
+        return
+
+    # FIX: Add try
+    parameters = dh.generate_parameters(generator=2, key_size=512)
+    skey = parameters.generate_private_key()
+    peer_pkey= serialization.load_pem_public_key(peer_pkey_file.read())
+    derived_key = None
+    if isinstance(peer_pkey, dh.DHPublicKey):
+        shared_key = skey.exchange(peer_pkey)
+        derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=b'handshake data').derive(shared_key)
+    else:
+        print("Error: Invalid key type")
+        return
+
+    message = file_r.read()
+    chacha = ChaCha20Poly1305(derived_key)
+    nonce = os.urandom(12)
+    ciphertext = chacha.encrypt(nonce, message, None)
+    file_w = open(f"{filename}.enc", 'wb')
+    nonce_ciphertext = mkpair(nonce, ciphertext)
+    to_send = mkpair(peer_pkey.public_bytes(Encoding.PEM, PublicFormat.SubjectPublicKeyInfo), nonce_ciphertext)
+    file_w.write(to_send)
+
+    file_r.close()
+    file_w.close()
+    skey_file.close()
+    peer_pkey_file.close()
+
+
+def decrypt(user,filename):
+    skey_file = open(f"{user}.sk", 'rb')
+    if skey_file is None:
+        print(f"No private key found for {user}")
+        return
+
+    file_r = open(filename, 'rb')
+    if file_r is None:
+        print(f"File {filename} not found")
+        return
+
+    message_bytes = file_r.read()
+    peer_pkey_bytes, ciphertext = unpair(message_bytes)
+    # FIX: Add try
+    skey  = serialization.load_pem_private_key(skey_file.read(), None)
+    peer_pkey = serialization.load_pem_public_key(peer_pkey_bytes)
+    derived_key = None
+    if isinstance(skey, dh.DHPrivateKey) and isinstance(peer_pkey, dh.DHPublicKey):
+        shared_key = skey.exchange(peer_pkey)
+        derived_key = HKDF(algorithm=hashes.SHA256(), length=32, salt=None, info=b'handshake data').derive(shared_key)
+    else:
+        print("Error: Invalid key type")
+        return
+
+    nonce, encrypted_message = unpair(ciphertext)
+    chacha = ChaCha20Poly1305(derived_key)
+    message = chacha.decrypt(nonce, encrypted_message, None)
+    file_w = open(f"{filename}.decrypt", 'wb')
+    file_w.write(message)
+    file_r.close()
+    file_w.close()
+    skey_file.close()
+
+
+
+
+def mkpair(x, y):
+    """ produz uma byte-string contendo o tuplo '(x,y)' ('x' e 'y' são byte-strings) """
+    len_x = len(x)
+    len_x_bytes = len_x.to_bytes(2, 'little')
+    return len_x_bytes + x + y
+
+def unpair(xy):
+    """ extrai componentes de um par codificado com 'mkpair' """
+    len_x = int.from_bytes(xy[:2], 'little')
+    x = xy[2:len_x+2]
+    y = xy[len_x+2:]
+    return x, y
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Program to perform operations using ChaCha20 cipher on files",
+    )
+
+    subparsers = parser.add_subparsers(dest="operation", help="Operation to perform")
+
+    # Setup subcommand
+    setup_parser = subparsers.add_parser("setup", help="Setup diffie helman key pair for user")
+    setup_parser.add_argument("user", help="User for which to setup the key pair")
+
+    # Encrypt subcommand
+    enc_parser = subparsers.add_parser("enc", help="Encrypt a file")
+    enc_parser.add_argument("user", help="User for which to encrypt the file")
+    enc_parser.add_argument("file", help="File to encrypt")
+
+    # Decrypt subcommand
+    decrypt_parser = subparsers.add_parser("dec", help="Decrypt a file")
+    decrypt_parser.add_argument("user", help="Target user for the decryption")
+    decrypt_parser.add_argument("file", help="File to be decrypted")
+
+    args = parser.parse_args()
+    match args.operation:
+        case "setup":
+            user = args.user
+            setup(user)
+        case "enc":
+            user = args.user
+            file = args.file
+            encrypt(user,file)
+        case "dec":
+            user = args.user
+            file = args.file
+            decrypt(user,file)
+
+if __name__ == "__main__":
+    main()