1import Pyro4
2import Pyro4.errors
3from diffiehellman import DiffieHellman
4
5Pyro4.config.SERVERTYPE = "multiplex"
6
7
8@Pyro4.expose
9class SecretStuff(object):
10    def process(self, message):
11        if not self._pyroDaemon._pyroHmacKey:
12            raise Pyro4.errors.PyroError("no hmac key has been set, can't call this method!")
13        print("Got the message:", message)
14        return "message was received ok"
15
16ns = Pyro4.locateNS()
17daemon = Pyro4.Daemon()
18daemon._pyroHmacKey = b"will be set to shared secret key by KeyExchange"
19uri = daemon.register(SecretStuff)
20ns.register("example.dh.secretstuff", uri)
21
22
23@Pyro4.behavior(instance_mode="session")
24class KeyExchange(object):
25    def __init__(self):
26        print("New KeyExchange, initializing Diffie-Hellman")
27        self.dh = DiffieHellman(group=14)
28
29    @Pyro4.expose
30    def exchange_key(self, other_public_key):
31        print("received a public key, calculating shared secret...")
32        self.dh.make_shared_secret_and_key(other_public_key)
33        print("setting new shared secret key.")
34        global daemon
35        daemon._pyroHmacKey = self.dh.key
36        return self.dh.public_key
37
38
39# The key exchange service can't be part of the same daemon as the
40# other service because it must not have a Hmac key set on the daemon.
41# So we create another daemon without hmac key and combine it.
42key_daemon = Pyro4.Daemon()
43uri = key_daemon.register(KeyExchange)
44ns.register("example.dh.keyexchange", uri)
45ns._pyroRelease()
46
47key_daemon.combine(daemon)
48print("Starting server loop...")
49key_daemon.requestLoop()
50