Skip to content

Instantly share code, notes, and snippets.

@vankesteren
Last active October 13, 2025 07:36
Show Gist options
  • Save vankesteren/19243bcdf6943349f231f0cb1de3e17c to your computer and use it in GitHub Desktop.
Save vankesteren/19243bcdf6943349f231f0cb1de3e17c to your computer and use it in GitHub Desktop.
Diffie-Hellman key exchange
from random import randint
class Agent:
"""Agent that can share their public key with another agent to compute a shared secret, without sharing that secret."""
def __init__(self, name: str, prime: int = 2**17 - 1, base: int = 2**5 - 1) -> None:
self.name = name
self.prime = prime
self.base = base
self.privkey = self.generate_private_key()
self.pubkey = self.generate_public_key()
self.sharedkey = None
def generate_private_key(self) -> int:
return randint(int(1e5), int(1e9 - 1))
def generate_public_key(self) -> int:
return pow(self.base, self.privkey, self.prime)
def compute_shared_key(self, other_pubkey: int) -> None:
self.sharedkey = pow(other_pubkey, self.privkey, self.prime)
def __repr__(self) -> str:
out = f"\n[ Agent: {self.name} ]"
out += f"\n Prime: {self.prime}"
out += f"\n Base: {self.base}"
out += f"\n🔐 Private key: {self.privkey}"
out += f"\n🔑 Public key: {self.pubkey}"
if self.sharedkey is not None:
out += f"\n🗝️ Shared key: {self.sharedkey}"
return out + "\n"
def main():
print("Creating Alice and Bob.")
alice = Agent("Alice")
bob = Agent("Bob")
print("\nHere they are:")
print(alice)
print(bob)
print("\nSharing pubkeys")
print("- Alice -> Bob")
bob.compute_shared_key(alice.pubkey)
print("- Bob -> Alice")
alice.compute_shared_key(bob.pubkey)
print("\nNow they both have the same shared key!")
print(alice)
print(bob)
assert alice.sharedkey == bob.sharedkey
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment