from stem.control import Controller
import stem.connection
import paho.mqtt.client as mqtt
import ssl
import socks
import socket
import json
import configparser
import argparse
from datetime import datetime
from os import environ

parser = argparse.ArgumentParser(description='Broadcast SSH hidden service hostname via MQTT')

parser.add_argument('--config-dir', nargs='?', dest='configPath', default='/etc/torch',
                   help='configuration directory (default: /etc/torch)')

args = parser.parse_args()

configPath = args.configPath

if "TORCH_CONFIG_DIR" in environ:
  configPath = environ.get("TORCH_CONFIG_DIR")

if not configPath.endswith("/"):
  configPath = configPath + "/"

print("Using torch configuration path: " + configPath)

config = configparser.ConfigParser()
config.read(configPath + "torch.conf")

torProxyPort = config['tor'].getint('ProxyPort', fallback = 9050)
torControllerPort = config['tor'].getint('ControllerPort', fallback = 9051)

sshPort = config['ssh'].getint('Port', fallback = 22)

mqttConfig = config['mqtt']
mqttBrokerHost = mqttConfig.get('BrokerHost', fallback = "localhost")
mqttBrokerPort = mqttConfig.getint('BrokerPort', fallback = 1883)
clientID = mqttConfig.get('ClientID', fallback = socket.gethostname())
mqttTopic = mqttConfig.get('Topic', fallback = "torch/%s/onion_url" % (clientID))

mqttRequireCertificate = mqttConfig.getboolean(
                                       'RequireCertificate', 
                                       fallback = False)

mqttCaFile = configPath + mqttConfig.get('CaFile')
mqttCertFile = configPath + mqttConfig.get('CertFile')
mqttKeyFile = configPath + mqttConfig.get('KeyFile')

with Controller.from_port(port = torControllerPort) as controller:
  
  protocolInfo = stem.connection.get_protocolinfo(controller)

  stem.connection.authenticate_safecookie(
                       controller, 
                       protocolInfo.cookie_path)
  
  print("Connected to Tor on port %s" % (torControllerPort))

  service = controller.create_ephemeral_hidden_service(
                             sshPort,
							 detached = True)

  onionAddress = "%s.onion" % (service.service_id)

  print("Created Tor Hidden Service for local port %s at %s" % (sshPort, onionAddress))

payload = { 
  'clientId': clientID, 
  'timestamp': datetime.now().strftime("%d-%b-%Y (%H:%M:%S.%f)"), 
  'onionAddress': onionAddress, 
  'sshPort': sshPort 
  }

client = mqtt.Client()
protocol = "mqtt"

if mqttRequireCertificate:
  client.tls_set(
          ca_certs = mqttCaFile, 
          certfile = mqttCertFile, 
          keyfile = mqttKeyFile, 
          cert_reqs=ssl.CERT_REQUIRED)
  protocol = "mqtts"

if mqttBrokerHost.endswith(".onion"):
  client.proxy_set(proxy_type=socks.SOCKS5, proxy_addr="localhost", proxy_port=torProxyPort)
  client.tls_insecure_set(True)

client.connect(mqttBrokerHost, mqttBrokerPort, 60)
client.publish(mqttTopic, json.dumps(payload))
print("Connected to MQTT Broker at %s://%s:%s/%s" % (protocol, mqttBrokerHost, mqttBrokerPort, mqttTopic))
print("Published payload: " + json.dumps(payload))

client.disconnect()
print("Disconnected from MQTT Broker")