From 7d764cb043a2600f912eb5d3a4a528aa45de974d Mon Sep 17 00:00:00 2001
From: Benjamin Dweck <bjdweck@gmail.com>
Date: Thu, 22 Oct 2020 11:07:26 +0200
Subject: [PATCH] Moved over to using mqtt.publish.single() instead of manual
 connection

---
 torch_agent/torch_agent.py | 70 ++++++++++++++++++--------------------
 1 file changed, 34 insertions(+), 36 deletions(-)

diff --git a/torch_agent/torch_agent.py b/torch_agent/torch_agent.py
index 831c185..70d01de 100755
--- a/torch_agent/torch_agent.py
+++ b/torch_agent/torch_agent.py
@@ -1,15 +1,16 @@
-from stem.control import Controller
-import stem.connection
-import paho.mqtt.client as mqtt
-import socks
-import ssl
-import socket
-import json
-import configparser
 import argparse
+import configparser
+import json
+import socket
+import ssl
 from datetime import datetime
 from os import environ
 
+import paho.mqtt.publish as publish
+import socks
+import stem.connection
+from stem.control import Controller
+
 
 def main():
     parser = argparse.ArgumentParser(description='Broadcast SSH hidden service hostname via MQTT')
@@ -45,8 +46,8 @@ def main():
     mqtt_broker_host = mqtt_config.get('BrokerHost', fallback="localhost")
     mqtt_broker_port = mqtt_config.getint('BrokerPort', fallback=1883)
     mqtt_broker_using_tor = mqtt_broker_host.endswith(".onion")
-    client_id = mqtt_config.get('ClientID', fallback=socket.gethostname())
-    mqtt_topic = mqtt_config.get('Topic', fallback="torch/%s/onion_url" % client_id)
+    mqtt_client_id = mqtt_config.get('ClientID', fallback=socket.gethostname())
+    mqtt_topic = mqtt_config.get('Topic', fallback="torch/%s/onion_url" % mqtt_client_id)
 
     mqtt_require_certificate = mqtt_config.getboolean(
         'RequireCertificate',
@@ -72,45 +73,42 @@ def main():
 
         print("Created Tor Hidden Service for local service on port %s at %s" % (ssh_port, onion_address))
 
-    def on_publish(pub_client):
-        print("Published payload!")
-        pub_client.disconnect()
-
-    def on_disconnect(pub_client):
-        print("Disconnected!")
-        pub_client.loop_stop()
-
-    client = mqtt.Client()
-    client.on_publish = on_publish
-    client.on_disconnect = on_disconnect
-    client.loop_start()
-
     protocol = "mqtt"
+    tls_args = None
+    proxy_args = None
 
     if mqtt_broker_using_tor:
-        client.proxy_set(proxy_type=socks.SOCKS5, proxy_addr=tor_proxy_host, proxy_port=tor_proxy_port)
+        proxy_args = {
+            'proxy_type': socks.SOCKS5,
+            'proxy_addr': tor_proxy_host,
+            'proxy_port': tor_proxy_port
+        }
     else:
         if mqtt_require_certificate:
             protocol = "mqtts"
-            client.tls_set(
-                ca_certs=mqtt_ca_file,
-                certfile=mqtt_cert_file,
-                keyfile=mqtt_key_file,
-                cert_reqs=ssl.CERT_REQUIRED)
+            tls_args = {
+                'ca_certs': mqtt_ca_file,
+                'certfile': mqtt_cert_file,
+                'keyfile': mqtt_key_file,
+                'cert_reqs': ssl.CERT_REQUIRED
+            }
 
-    print("Connecting to MQTT broker: %s://%s:%s/%s" % (protocol, mqtt_broker_host, mqtt_broker_port, mqtt_topic))
+    print("Publishing to MQTT broker: %s://%s:%s/%s" % (protocol, mqtt_broker_host, mqtt_broker_port, mqtt_topic))
     if mqtt_broker_using_tor:
         print("--> Using TOR proxy: %s:%s" % (tor_proxy_host, tor_proxy_port))
 
-    client.connect(mqtt_broker_host, mqtt_broker_port, 60)
-    print("Connected to MQTT Broker")
-
     payload = json.dumps({
-        'clientId': client_id,
+        'clientId': mqtt_client_id,
         'timestamp': datetime.now().strftime("%d-%b-%Y (%H:%M:%S.%f)"),
         'onionAddress': onion_address,
         'sshPort': ssh_port
     })
 
-    print("Publishing payload: " + payload)
-    client.publish(mqtt_topic, payload, qos=1)
+    publish.single(mqtt_topic,
+                   payload,
+                   qos=1,
+                   hostname=mqtt_broker_host,
+                   port=mqtt_broker_port,
+                   client_id=mqtt_client_id,
+                   tls=tls_args,
+                   proxy_args=proxy_args)