python书写socket Diffie-Hellman密钥交换算法

下面是使用Python编写的Socket Diffie-Hellman密钥交换算法的示例代码：

import socket
import math
import random

# 生成大质数p和原根g
def generate_prime_and_primitive_root():
    while True:
        p = random.randint(100, 1000)  # 生成随机数p
        if is_prime(p):
            break

    for i in range(2, p):
        if math.gcd(i, p-1) == 1:
            g = i
            return p, g

# 判断是否为质数
def is_prime(num):
    if num == 2 or num == 3:
        return True
    if num % 2 == 0 or num < 2:
        return False
    for i in range(3, int(num ** 0.5) + 1, 2):
        if num % i == 0:
            return False
    return True

# Diffie-Hellman密钥交换
def diffie_hellman(sock):
    # 生成大质数p和原根g
    p, g = generate_prime_and_primitive_root()

    # 发送p和g给对方
    sock.sendall(str(p).encode())
    sock.sendall(str(g).encode())

    # 生成私钥a并计算公钥A
    a = random.randint(2, p-2)
    A = pow(g, a, p)

    # 接收对方的公钥B
    B = int(sock.recv(1024).decode())

    # 计算共享密钥
    K = pow(B, a, p)

    # 发送公钥A给对方
    sock.sendall(str(A).encode())

    return K

# 主函数
def main():
    # 创建套接字并连接服务器
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    host = '127.0.0.1'  # 服务器IP地址
    port = 8000  # 服务器端口号
    sock.connect((host, port))

    # Diffie-Hellman密钥交换
    K = diffie_hellman(sock)

    # 输出共享密钥
    print('共享密钥K为：', K)

    # 关闭套接字
    sock.close()

if __name__ == '__main__':
    main()

在上面的代码中，我们首先使用`generate_prime_and_primitive_root()`函数生成了大质数p和原根g，并将它们发送给对方。然后，我们生成私钥a并计算公钥A，接收对方的公钥B，计算共享密钥K，最后将公钥A发送给对方。在`main()`函数中，我们创建套接字并连接服务器，执行Diffie-Hellman密钥交换，输出共享密钥K，最后关闭套接字。