Add P2P file sharing algorithm in Python

Author: Nikhil-karoriya

peer2peer_file_sharing/crypto/aes_crypto.py

@@ -0,0 +1,37 @@
+import zlib
+from cryptography.fernet import Fernet
+
+SKIP_COMPRESSION_EXTENSIONS = {
+    '.zip', '.gz', '.bz2', '.xz', '.rar', '.7z',
+    '.jpg', '.jpeg', '.png', '.gif', '.webp',
+    '.mp3', '.mp4', '.avi', '.mkv', '.mov',
+    '.pdf'
+}
+
+def generate_aes_key():
+    return Fernet.generate_key()
+
+def encrypt_chunk_with_aes(chunk: bytes, aes_key: bytes, file_extension: str) -> bytes:
+    fernet = Fernet(aes_key)
+    if file_extension.lower() not in SKIP_COMPRESSION_EXTENSIONS:
+        try:
+            chunk = zlib.compress(chunk)
+        except Exception as e:
+            print(f"\n[!] Compression failed: {e}")
+            
+    return fernet.encrypt(chunk)
+
+def decrypt_chunk_with_aes(chunk: bytes, aes_key: bytes, file_extension: str) -> bytes:
+    fernet = Fernet(aes_key)
+    try:
+        decrypted_data = fernet.decrypt(chunk)
+    except Exception as e:
+        raise Exception(f"\n[!] AES decryption failed: {e}")
+
+    if file_extension.lower() not in SKIP_COMPRESSION_EXTENSIONS:
+        try:
+            decrypted_data = zlib.decompress(decrypted_data)
+        except zlib.error:
+            print("\n[!] Warning: Failed to decompress — file may not be compressed")
+
+    return decrypted_data

peer2peer_file_sharing/crypto/rsa_crypto.py

@@ -0,0 +1,63 @@
+from cryptography.hazmat.primitives import serialization, hashes
+from cryptography.hazmat.primitives.asymmetric import rsa, padding
+
+def generate_rsa_key_pair(private_path: str, public_path: str):
+    private_key = rsa.generate_private_key(public_exponent=65537, key_size=2048)
+    
+    with open(private_path, "wb") as f:
+        
+        f.write(private_key.private_bytes(
+            serialization.Encoding.PEM,
+            serialization.PrivateFormat.TraditionalOpenSSL,
+            serialization.NoEncryption()
+        ))
+    
+    with open(public_path, "wb") as f:
+        
+        f.write(private_key.public_key().public_bytes(
+            serialization.Encoding.PEM,
+            serialization.PublicFormat.SubjectPublicKeyInfo
+        ))
+
+def encrypt_with_rsa_public_key(data: bytes, public_key_path: str) -> bytes:
+
+    with open(public_key_path, "rb") as f:
+        public_key = serialization.load_pem_public_key(f.read())
+
+    if not isinstance(public_key, rsa.RSAPublicKey):
+        raise TypeError("\n[!] The loaded public key is not an RSA key.")
+    
+    return public_key.encrypt(
+        data,
+        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
+    )
+
+def decrypt_with_rsa_private_key(ciphertext: bytes, private_key_path: str) -> bytes:
+    
+    with open(private_key_path, "rb") as f:
+        private_key = serialization.load_pem_private_key(f.read(), password=None)
+
+    if not isinstance(private_key, rsa.RSAPrivateKey):
+        raise TypeError("\n[!] The loaded private key is not an RSA key.")
+    
+    return private_key.decrypt(
+        ciphertext,
+        padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
+    )
+
+def is_valid_pem_key(file_path: str, is_private: bool = False) -> bool:
+    
+    try:
+        with open(file_path, "rb") as f:
+            data = f.read()
+            if not data:
+                return False
+            if is_private:
+                serialization.load_pem_private_key(data, password=None)
+            else:
+                serialization.load_pem_public_key(data)
+        return True
+    
+    except Exception:
+        return False
+    

peer2peer_file_sharing/peer.py

@@ -0,0 +1,191 @@
+"""
+P2P File Sharing Implementation
+Author: Nikhil Karoriya
+Description: This module implements a secure peer-to-peer file sharing algorithm
+using socket programming, AES for data encryption, RSA for key exchange and Zlib for data compression
+
+Usage:
+    python peer.py --listen-port 5001
+
+You will be prompted for:
+
+    Send file (y/n)? y
+    Enter file path to send: sample.txt
+    Enter receiver's IP address: localhost (or receiver's IP address)
+    Enter receiver's port: 6001
+"""
+
+
+import os
+import socket
+import threading
+import traceback
+from tqdm import tqdm
+import argparse
+from time import sleep
+from crypto.rsa_crypto import (
+    decrypt_with_rsa_private_key,
+    encrypt_with_rsa_public_key,
+    generate_rsa_key_pair,
+    is_valid_pem_key
+)
+from crypto.aes_crypto import generate_aes_key, encrypt_chunk_with_aes, decrypt_chunk_with_aes
+from utils.file_utils import ensure_dir, sha256_digest_stream, is_valid_ip, is_valid_port
+from utils.file_utils import CHUNK_SIZE
+
+PRIVATE_KEY_PATH = 'keys/private_key.pem'
+PUBLIC_KEY_PATH = 'keys/public_keys.pem'
+RECEIVE_DIR = 'received_files'
+
+ensure_dir(RECEIVE_DIR)
+ensure_dir("keys")
+
+if not is_valid_pem_key(PRIVATE_KEY_PATH, is_private=True) or not is_valid_pem_key(PUBLIC_KEY_PATH, is_private=False):
+    print("\n[!] RSA key missing or invalid. Regenerating...")
+    generate_rsa_key_pair(PRIVATE_KEY_PATH, PUBLIC_KEY_PATH)
+    print("\n[+] RSA key pair regenerated.")
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--listen-port", type=int, required=True)
+args = parser.parse_args()
+
+LISTEN_PORT = args.listen_port
+
+def peer_listener():
+    try:
+        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server_socket.bind(('0.0.0.0', LISTEN_PORT))
+        server_socket.listen(1)
+        print(f"\n[+] Listening on port {LISTEN_PORT}...")
+
+        while True:
+            client_socket, addr = server_socket.accept()
+            print(f"\n\n[+] Incoming connection from {addr}\n")
+
+            try:
+                key_size = int.from_bytes(client_socket.recv(4), 'big')
+                encrypted_key = client_socket.recv(key_size)
+
+                name_len = int.from_bytes(client_socket.recv(4), 'big')
+                file_name = client_socket.recv(name_len).decode()
+
+                extension = int.from_bytes(client_socket.recv(4), 'big')
+                file_extension = client_socket.recv(extension).decode()
+
+                total_size = int.from_bytes(client_socket.recv(8), 'big')
+                
+                aes_key = decrypt_with_rsa_private_key(encrypted_key, PRIVATE_KEY_PATH)
+
+                file_path = os.path.join(RECEIVE_DIR, file_name)
+                
+                with open(file_path, 'wb') as f_out, tqdm(total=total_size, desc=f"[+] Receiving {file_name}", unit="B", unit_scale=True) as pbar:
+                    received_bytes = 0
+                    while received_bytes < total_size:
+                        rcv_chunk_size = int.from_bytes(client_socket.recv(4), 'big')
+                        encrypted_chunk = b""
+                        
+                        while len(encrypted_chunk) < rcv_chunk_size:
+                            part = client_socket.recv(rcv_chunk_size - len(encrypted_chunk))
+                            if not part:
+                                raise Exception("\n[!] Connection lost during transfer.")
+                            encrypted_chunk += part
+
+                        decrypted_chunk = decrypt_chunk_with_aes(encrypted_chunk, aes_key, file_extension)
+                        f_out.write(decrypted_chunk)
+                        received_bytes += len(decrypted_chunk)
+                        pbar.update(len(decrypted_chunk))
+
+                print(f"\n[+] File saved to {file_path}")
+                file_hash = sha256_digest_stream(file_path)
+                print(f"\n[+] SHA256 Hash: {file_hash}")
+
+            except Exception as e:
+                print(f"\n[!] ERROR receiving file: {str(e)}\n{traceback.format_exc()}")
+            finally:
+                client_socket.close()
+
+    except Exception as e:
+        print(f"\n[!] Server failed: {str(e)}\n{traceback.format_exc()}")
+
+def send_file(file_path, peer_ip, peer_port):
+    try:
+        
+        file_name = os.path.basename(file_path)
+        file_extension = os.path.splitext(file_name)[1].lower()
+
+        aes_key = generate_aes_key()
+        encrypted_key = encrypt_with_rsa_public_key(aes_key, PUBLIC_KEY_PATH)
+        file_size = os.path.getsize(file_path)
+
+        client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        client_socket.connect((peer_ip, peer_port))
+
+        client_socket.sendall(len(encrypted_key).to_bytes(4, 'big'))
+        client_socket.sendall(encrypted_key)
+
+        client_socket.sendall(len(file_name.encode()).to_bytes(4, 'big'))
+        client_socket.sendall(file_name.encode())
+
+        client_socket.sendall(len(file_extension.encode()).to_bytes(4, 'big'))
+        client_socket.sendall(file_extension.encode())
+
+        client_socket.sendall(file_size.to_bytes(8, 'big'))
+        
+        print()
+
+        with open(file_path, 'rb') as f, tqdm(total=file_size, desc=f"[+] Sending {file_name}", unit="B", unit_scale=True) as pbar:
+            for chunk in iter(lambda: f.read(CHUNK_SIZE), b''):
+                encrypted_chunk = encrypt_chunk_with_aes(chunk, aes_key, file_extension)
+                client_socket.sendall(len(encrypted_chunk).to_bytes(4, 'big'))
+                client_socket.sendall(encrypted_chunk)
+                pbar.update(len(chunk))
+
+        print(f"\n[+] File '{file_name}' sent successfully to {peer_ip}:{peer_port}.")
+        client_socket.close()
+
+    except Exception as e:
+        print(f"\n[!] Failed to send file: {str(e)}\n{traceback.format_exc()}")
+
+
+if __name__ == '__main__':
+    
+    try:
+        threading.Thread(target=peer_listener, daemon=True).start()
+        sleep(0.5)
+
+        while True:
+            user_input = input("\nSend file [y], wait [w], or exit [e]? ").strip().lower()
+
+            if user_input == 'e':
+                print("\n[INFO] Exiting...")
+                break
+
+            elif user_input == 'y':
+                file_path = input("\nEnter file path to send: ").strip()
+                peer_ip = input("\nEnter receiver's IP address: ").strip()
+                peer_port_input = input("\nEnter receiver's port: ").strip()
+
+                if not is_valid_port(peer_port_input) or not is_valid_ip(peer_ip):
+                    print("\n[!] Receiver IP or port not correct/specified.")
+                    continue
+
+                if not os.path.isfile(file_path):
+                    print("\n[!] File does not exist.")
+                    continue
+
+                peer_port = int(peer_port_input)
+                send_file(file_path, peer_ip, peer_port)
+
+            elif user_input == 'w':
+                print("\n[INFO] Waiting for incoming transfers...")
+
+            else:
+                print("\n[!] Invalid input. Use 'y', 'w', or 'e'.")
+
+    except KeyboardInterrupt:
+        print("\n[INFO] Exiting by keyboard interrupt.")
+        
+    except Exception as e:
+        print(f"\n[!] An error occurred: {str(e)}\n{traceback.format_exc()}")
+
+        

peer2peer_file_sharing/utils/file_utils.py

@@ -0,0 +1,36 @@
+import os
+import hashlib
+import ipaddress
+
+CHUNK_SIZE = 64 * 1024
+
+def sha256_digest_stream(path: str) -> str:
+
+    hasher = hashlib.sha256()
+    with open(path, 'rb') as f:
+        for chunk in iter(lambda: f.read(CHUNK_SIZE), b''):
+            hasher.update(chunk)
+
+    return hasher.hexdigest()
+
+def ensure_dir(directory: str):
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+
+def is_valid_ip(ip_str):
+    try:
+        if ip_str.strip().lower() == "localhost":
+            return True
+        ipaddress.ip_address(ip_str)
+        return True
+    
+    except ValueError:
+        return False
+
+def is_valid_port(port_str):
+    try:
+        port = int(port_str)
+        return 1 <= port <= 65535
+    
+    except ValueError:
+        return False