# 加密傳訊息聊天程式
此為我學習玩AES和RSA後生成的實作專題

from Crypto.PublicKey import RSA
from Crypto.Cipher import AES,PKCS1_OAEP
import socket
import threading
import base64
import os

#產生RSA密鑰
def generate_RSA_key():
    key=RSA.generate(2048) # 產生 2048-bit RSA 密鑰
    private_key=key.export_key() 
    public_key=key.publickey().export_key() 
    return private_key, public_key

#AES 加密
def encrypt_message_AES(key, plaintext):
    cipher=AES.new(key,AES.MODE_EAX)
    ciphertext, tag=cipher.encrypt_and_digest(plaintext.encode())
    return base64.b64encode(cipher.nonce+tag+ciphertext).decode()

#AES 解密
def decrypt_message_AES(key,encrypted_text):
    raw=base64.b64decode(encrypted_text)
    nonce,tag,ciphertext=raw[:16],raw[16:32],raw[32:]
    cipher=AES.new(key, AES.MODE_EAX, nonce=nonce)
    return cipher.decrypt_and_verify(ciphertext, tag).decode()
# RSA 加密
def encrypt_aes_key_rsa(public_key, aes_key):
    recipient_key = RSA.import_key(public_key)
    cipher_rsa = PKCS1_OAEP.new(recipient_key)
    encrypted_key = cipher_rsa.encrypt(aes_key)
    return base64.b64encode(encrypted_key).decode()

# RSA 解密
def decrypt_aes_key_rsa(private_key, encrypted_key):
    private_rsa = RSA.import_key(private_key)
    cipher_rsa = PKCS1_OAEP.new(private_rsa)
    return cipher_rsa.decrypt(base64.b64decode(encrypted_key))

# Server 端（接收並解密訊息）
def server():
    private_key, public_key = generate_RSA_key() # 產生 RSA 金鑰
    aes_key = os.urandom(16)  # 產生 AES 密鑰
    
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.bind(('127.0.0.1', 65432))
    server_socket.listen(1)
    print("[SERVER] 等待連線...")
    
    conn, addr = server_socket.accept()
    print(f"[SERVER] 來自 {addr} 的連線")
    conn.send(public_key)  # 傳送公鑰
    
    encrypted_aes_key = conn.recv(1024).decode()
    aes_key = decrypt_aes_key_rsa(private_key, encrypted_aes_key)  # 解密 AES 密鑰
    print("[SERVER] AES 金鑰已解密！")
    
    while True:
        encrypted_msg = conn.recv(1024).decode()
        if not encrypted_msg:
            break
        print(f"[SERVER] 接收到加密訊息: {encrypted_msg}")
        decrypted_msg = decrypt_message_AES(aes_key, encrypted_msg)  # AES 解密訊息
        print(f"[SERVER] 客戶端: {decrypted_msg}")
        
        response = input("[SERVER] 輸入回覆: ")
        conn.send(encrypt_message_AES(aes_key, response).encode()) # 回覆訊息並加密

# Client 端（加密並發送訊息）
def client():
    client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    client_socket.connect(('127.0.0.1', 65432))
    
    public_key = client_socket.recv(1024)  # 接收 Server 公鑰
    aes_key = os.urandom(16)  # 產生 AES 金鑰
    encrypted_aes_key = encrypt_aes_key_rsa(public_key, aes_key)
    client_socket.send(encrypted_aes_key.encode())  # 發送加密後的 AES 金鑰
    
    print("[CLIENT] 已加密 AES 金鑰並傳送！")
    
    while True:
        msg = input("[CLIENT] 輸入訊息: ")
        encrypted_msg = encrypt_message_AES(aes_key, msg)  # AES 加密訊息
        print(f"[CLIENT] 加密後的訊息: {encrypted_msg}")
        client_socket.send(encrypted_msg.encode())
        
        response = client_socket.recv(1024).decode()
        print(f"[CLIENT] 伺服器: {decrypt_message_AES(aes_key, response)}") # 解密回覆訊息

# 啟動伺服器或客戶端
if __name__ == "__main__":
    choice = input("啟動伺服器 (s) 或 客戶端 (c)？").strip().lower()
    if choice == 's':
        server()
    elif choice == 'c':
        client()
    else:
        print("請選擇 's' 或 'c'")