server.py

import base64
import os
import socket
import sys
import threading

from libs import aes, path, web_io


def has_key(filename: str) -> bool:
    file_path = path.get_key_path(filename)
    return os.path.exists(file_path)


def main_recv(s: socket.socket, key: bytes, stop_event: threading.Event):
    """副线程用于接受另一个用户的消息"""
    while True:
        try:
            en_data = web_io.recv_msg(s)
            data = aes.decrypt_ecb(key, en_data).decode("utf-8")
            print(data)
        except web_io.SocketUtilsError:
            print("远程主机强迫关闭了当前的连接")
            stop_event.set()
            break


def run():
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("127.0.0.1", 12345))
        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        s.listen(1)
        conn, addr = s.accept()

        rsa_keys = ["pubkey.pem", "privkey.pem"]
        if not all(map(has_key, rsa_keys)):
            _ = sys.stderr.write("RSA keys is uncompletely!\n")
            sys.exit(-1)
        content = []
        with open(f"{path.get_key_path('pubkey.pem')}", "r") as f:
            content = f.readlines()
        d = content[1].encode("utf-8")
        e = content[2].encode("utf-8")
        pubkey = (d + e).strip()

        # 发送 rsa 公钥
        with conn:
            web_io.send_msg(pubkey, conn)

            # 接受 aes 密钥
            aes_key = web_io.recv_msg(conn)
            aes_key = base64.b64decode(aes_key)

            # 使用密钥加密通讯
            stop_event = threading.Event()
            t = threading.Thread(
                target=main_recv,
                args=(
                    conn,
                    aes_key,
                    stop_event,
                ),
                daemon=True,
            )
            t.start()
            while True:
                msg = input("> ")
                if msg == "quit" or stop_event.is_set():
                    break
                en_msg = aes.encrypt_ecb(aes_key, msg.encode("utf-8"))
                web_io.send_msg(en_msg, conn)


if __name__ == "__main__":
    run()