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)