secure_chat.py

import os
import json
import base64
import time
import sys
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from cryptography.exceptions import InvalidTag

# ========== 目录处理 ==========
def get_script_dir():
    if getattr(sys, 'frozen', False):
        return os.path.dirname(sys.executable)
    return os.path.dirname(os.path.abspath(__file__))

SCRIPT_DIR = get_script_dir()
def local_path(filename): return os.path.join(SCRIPT_DIR, filename)

PRIVATE_KEY_PATH = local_path("private_key.pem")
PUBLIC_KEY_PATH = local_path("public_key.pem")

# ========== 密钥处理 ==========
def generate_key_pair():
    private_key = rsa.generate_private_key(public_exponent=65537, key_size=4096)
    public_key = private_key.public_key()

    with open(PRIVATE_KEY_PATH, "wb") as f:
        f.write(private_key.private_bytes(
            serialization.Encoding.PEM,
            serialization.PrivateFormat.TraditionalOpenSSL,
            serialization.NoEncryption()
        ))
    with open(PUBLIC_KEY_PATH, "wb") as f:
        f.write(public_key.public_bytes(
            serialization.Encoding.PEM,
            serialization.PublicFormat.SubjectPublicKeyInfo
        ))
    print("[+] 密钥对生成成功")

def load_public_key():
    if not os.path.exists(PUBLIC_KEY_PATH):
        raise FileNotFoundError("缺少 public_key.pem，请先生成密钥对（选项 1）")
    with open(PUBLIC_KEY_PATH, "rb") as f:
        return serialization.load_pem_public_key(f.read())

def load_private_key():
    if not os.path.exists(PRIVATE_KEY_PATH):
        raise FileNotFoundError("缺少 private_key.pem，请先生成密钥对（选项 1）")
    with open(PRIVATE_KEY_PATH, "rb") as f:
        return serialization.load_pem_private_key(f.read(), password=None)

# ========== 加密文本 ==========
def encrypt_text():
    try:
        public_key = load_public_key()
    except FileNotFoundError as e:
        print("[-]", e)
        return

    print("请输入要加密的文本（以 END 结束）：")
    lines = []
    while True:
        line = input()
        if line.strip().upper() == "END":
            break
        lines.append(line)
    plaintext = "\n".join(lines).encode()

    aes_key, iv = os.urandom(32), os.urandom(12)
    encryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv)).encryptor()
    ciphertext = encryptor.update(plaintext) + encryptor.finalize()
    tag = encryptor.tag

    encrypted_key = public_key.encrypt(
        aes_key,
        padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )

    data = {
        "enc_key": base64.b64encode(encrypted_key).decode(),
        "iv": base64.b64encode(iv).decode(),
        "tag": base64.b64encode(tag).decode(),
        "ciphertext": base64.b64encode(ciphertext).decode()
    }

    filename = local_path(f"{int(time.time())}.json")
    with open(filename, "w") as f:
        json.dump(data, f, indent=2)
    print(f"[+] 文本加密完成，保存为 {os.path.basename(filename)}")

# ========== 解密文本 ==========
def decrypt_text(path):
    try:
        with open(os.path.abspath(path), "r") as f:
            data = json.load(f)
        enc_key = base64.b64decode(data["enc_key"])
        iv = base64.b64decode(data["iv"])
        tag = base64.b64decode(data["tag"])
        ciphertext = base64.b64decode(data["ciphertext"])
    except Exception:
        print("[-] 无效或损坏的 JSON 文件")
        return

    try:
        private_key = load_private_key()
    except FileNotFoundError as e:
        print("[-]", e)
        return

    try:
        aes_key = private_key.decrypt(
            enc_key,
            padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
        )
        decryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv, tag)).decryptor()
        plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        print("\n[+] 解密成功：\n")
        print(plaintext.decode())
    except InvalidTag:
        print("[-] 解密失败，数据完整性校验错误")
    except Exception as e:
        print("[-] 解密失败：", e)

# ========== 加密文件 ==========
def encrypt_file():
    try:
        public_key = load_public_key()
    except FileNotFoundError as e:
        print("[-]", e)
        return

    path = input("请输入要加密的文件路径：").strip()
    if not os.path.isfile(path):
        print("[-] 文件不存在")
        return

    with open(path, "rb") as f:
        file_data = f.read()

    aes_key, iv = os.urandom(32), os.urandom(12)
    encryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv)).encryptor()
    ciphertext = encryptor.update(file_data) + encryptor.finalize()
    tag = encryptor.tag

    encrypted_key = public_key.encrypt(
        aes_key,
        padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )

    timestamp = int(time.time())
    out_file = local_path(f"{timestamp}.enc")
    with open(out_file, "wb") as f:
        f.write(b"ENCFILE\n")
        f.write(base64.b64encode(encrypted_key) + b"\n")
        f.write(base64.b64encode(iv) + b"\n")
        f.write(base64.b64encode(tag) + b"\n")
        f.write(base64.b64encode(ciphertext))

    print(f"[+] 文件加密完成，保存为 {os.path.basename(out_file)}")

# ========== 解密文件 ==========
def decrypt_file(path):
    try:
        with open(os.path.abspath(path), "rb") as f:
            lines = f.read().split(b"\n")
        if lines[0] != b"ENCFILE":
            print("[-] 无效文件格式")
            return

        enc_key = base64.b64decode(lines[1])
        iv = base64.b64decode(lines[2])
        tag = base64.b64decode(lines[3])
        ciphertext = base64.b64decode(b"\n".join(lines[4:]))
    except Exception:
        print("[-] 文件读取失败")
        return

    try:
        private_key = load_private_key()
    except FileNotFoundError as e:
        print("[-]", e)
        return

    try:
        aes_key = private_key.decrypt(
            enc_key,
            padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
        )
        decryptor = Cipher(algorithms.AES(aes_key), modes.GCM(iv, tag)).decryptor()
        plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        output_path = local_path(f"decrypted_{int(time.time())}")
        with open(output_path, "wb") as f:
            f.write(plaintext)
        print(f"[+] 解密成功，保存为 {os.path.basename(output_path)}")
    except Exception as e:
        print("[-] 解密失败：", str(e))

# ========== 菜单 ==========
def print_help():
    print("""
======== 安全聊天脚本说明 ========

1. 生成密钥对（保存为 public_key.pem / private_key.pem）
2. 加密文本（输入内容 → 输出 timestamp.json）
3. 解密文本（输入 .json 文件 → 控制台输出）
4. 加密文件（输入路径 → 输出 timestamp.enc）
5. 解密文件（输入 .enc 文件 → 输出 decrypted_时间戳）
6. 显示帮助
7. 退出
""")

def menu():
    print_help()
    while True:
        choice = input("\n请选择操作（1-7）：").strip()
        if choice == "1": generate_key_pair()
        elif choice == "2": encrypt_text()
        elif choice == "3": decrypt_text(input("请输入 JSON 文件路径：").strip())
        elif choice == "4": encrypt_file()
        elif choice == "5": decrypt_file(input("请输入加密文件路径：").strip())
        elif choice == "6": print_help()
        elif choice == "7":
            print("退出程序。")
            break
        else:
            print("请输入有效选项（1-7）")

if __name__ == "__main__":
    try:
        menu()
    except KeyboardInterrupt:
        print("\n[!] 用户中断，程序退出")