From e580fd343ba6355d336a1733420b174b09fc10f7 Mon Sep 17 00:00:00 2001
From: Benjamin Dweck <bjdweck@gmail.com>
Date: Thu, 15 Oct 2020 19:07:17 +0200
Subject: [PATCH] Can receive message and record client record to file

---
 torch_sub/test/integration_pub_sub.py | 44 ++++++++++++++++++++-----
 torch_sub/torch_sub.py                | 47 ++++++++++++++++++++++-----
 2 files changed, 74 insertions(+), 17 deletions(-)

diff --git a/torch_sub/test/integration_pub_sub.py b/torch_sub/test/integration_pub_sub.py
index 5c87a48..981f2f4 100644
--- a/torch_sub/test/integration_pub_sub.py
+++ b/torch_sub/test/integration_pub_sub.py
@@ -1,5 +1,7 @@
 import json
 import ssl
+import threading
+import time
 import unittest
 from datetime import datetime
 
@@ -9,10 +11,15 @@ from torch_sub import torch_sub
 
 host = "mqtt.example.com"
 port = 8883
-config_path = "agent-config/"
-mqtt_ca_file = config_path + "ca.crt"
-mqtt_cert_file = config_path + "vagrant.crt"
-mqtt_key_file = config_path + "vagrant.key"
+agent_config_path = "agent-config/"
+mqtt_ca_file = agent_config_path + "ca.crt"
+mqtt_cert_file = agent_config_path + "vagrant.crt"
+mqtt_key_file = agent_config_path + "vagrant.key"
+
+subscriber_config_path = "subscriber-config/"
+subscriber_ca_file = subscriber_config_path + "ca.crt"
+subscriber_cert_file = subscriber_config_path + "subscriber.crt"
+subscriber_key_file = subscriber_config_path + "subscriber.key"
 
 
 def agent_connect():
@@ -48,19 +55,40 @@ def agent_publish(client_id, onion_hostname):
     pass
 
 
+def subscriber_thread(host, port, filename, topic, cafile, certfile, keyfile):
+    torch_sub.attach(host,
+                     port,
+                     filename,
+                     topic=topic,
+                     mqtt_ca_file=cafile,
+                     mqtt_cert_file=certfile,
+                     mqtt_key_file=keyfile)
+
+
 class GivenBrokerAndTorchAgent(unittest.TestCase):
     def test_when_agent_publishes_should_get_hostname_from_subscriber(self):
+        outfile = "clients.json"
 
-        clients_json = "clients.json"
+        threading.Thread(target=subscriber_thread,
+                         args=(host,
+                               port,
+                               outfile,
+                               "torch/+/wake",
+                               subscriber_ca_file,
+                               subscriber_cert_file,
+                               subscriber_key_file),
+                         daemon=True).start()
 
-        torch_sub.attach(host, port, clients_json)
+        time.sleep(0.5)
 
         agent_publish("client1", "crazyonion.onion")
 
-        file = open(clients_json, "r")
+        time.sleep(0.5)
+
+        file = open(outfile, "r")
         response = json.load(file)
 
-        self.assertEqual(response['results']['client1']['onionAddress'], "crazyonion.onion")
+        self.assertEqual(response['client1']['onionAddress'], "crazyonion.onion")
 
 
 if __name__ == '__main__':
diff --git a/torch_sub/torch_sub.py b/torch_sub/torch_sub.py
index fd9b873..553bd28 100644
--- a/torch_sub/torch_sub.py
+++ b/torch_sub/torch_sub.py
@@ -1,15 +1,44 @@
 import json
 
+import paho.mqtt.client as cl
+import paho.mqtt.subscribe as mqtt
 
-def attach(host, port, datafile):
+datafile = "clients.json"
 
-    response = {
-        'results': {
-            'client1': {
-                'onionAddress': 'crazytime'
-            }
-        }
-    }
+def updateOnion(client, userdata, message):
+
+    with open(datafile, 'r') as infile:
+        database = json.load(infile)
+
+    payload = message.payload.decode('utf-8')
+
+    print("Payload: %s" % (payload))
+
+    response = json.loads(payload)
+
+    print(response)
+    print("Response: %s" % (response))
+
+    print("Database: %s" % (database))
+    database[response['clientId']] = response
+
+    print("got one! %s %s %s" % (client, userdata, payload))
 
     with open(datafile, 'w') as outfile:
-        json.dump(response, outfile)
\ No newline at end of file
+        json.dump(database, outfile)
+
+
+
+def attach(host, port, datafileIn, mqtt_ca_file=None, mqtt_cert_file=None, mqtt_key_file=None, topic=None):
+
+    datafile = datafileIn
+
+    mqtt.callback(updateOnion,
+                  topic,
+                  hostname=host,
+                  port=port,
+                  tls={
+                      'ca_certs': mqtt_ca_file,
+                      'certfile': mqtt_cert_file,
+                      'keyfile': mqtt_key_file
+                  })
\ No newline at end of file