用python模拟diffie-hellman算法

好的，下面是Python代码实现Diffie-Hellman算法：

import random

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

def gen_prime():
    """生成一个大素数"""
    while True:
        num = random.randint(100, 1000)
        if is_prime(num):
            return num

def gen_primitive_root(p):
    """生成一个原根"""
    factors = []
    phi = p - 1
    n = phi
    for i in range(2, int(n ** 0.5) + 1):
        if n % i == 0:
            factors.append(i)
            while n % i == 0:
                n /= i
    if n > 1:
        factors.append(n)
    for g in range(2, p):
        flag = True
        for factor in factors:
            if pow(g, phi // factor, p) == 1:
                flag = False
                break
        if flag:
            return g

def gen_key(p, g):
    """生成密钥"""
    a = random.randint(2, p - 2)
    A = pow(g, a, p)
    return a, A

def encrypt(key, B, p):
    """加密"""
    return pow(B, key, p)

# 生成大素数和原根
p = gen_prime()
g = gen_primitive_root(p)

# 生成密钥
a, A = gen_key(p, g)
b, B = gen_key(p, g)

# 加密和解密
k1 = encrypt(a, B, p)
k2 = encrypt(b, A, p)

# 输出结果
print("p:", p)
print("g:", g)
print("a:", a, "A:", A)
print("b:", b, "B:", B)
print("k1:", k1)
print("k2:", k2)

代码中，我们首先实现了判断一个数是否为素数的函数 `is_prime()`，然后实现了生成一个大素数的函数 `gen_prime()` 和生成一个原根的函数 `gen_primitive_root()`。接着，我们实现了生成密钥的函数 `gen_key()`，该函数生成一个随机整数作为私钥，然后计算出公钥。最后，我们实现了加密和解密的函数 `encrypt()`，该函数接受一个密钥和对方的公钥，返回加密后的结果。

在主程序中，我们先生成了大素数和原根，然后生成了两个密钥，并分别用对方的公钥进行加密和解密，得到了最终的密钥。最后，我们输出了各个参数和最终的密钥。