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
from datetime import datetime
from os import environ


def main():
    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()

    config_path = args.configPath

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

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

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

    config = configparser.ConfigParser()

    configuration_file_path = config_path + "torch.conf"
    print("Reading configuration file at '%s'" % configuration_file_path)
    config.read(configuration_file_path)

    tor_proxy_host = config['tor'].get('ProxyHost', fallback="127.0.0.1")
    tor_proxy_port = config['tor'].getint('ProxyPort', fallback=9050)
    tor_controller_port = config['tor'].getint('ControllerPort', fallback=9051)

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

    mqtt_config = config['mqtt']
    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_require_certificate = mqtt_config.getboolean(
        'RequireCertificate',
        fallback=False)

    mqtt_ca_file = config_path + mqtt_config.get('CaFile')
    mqtt_cert_file = config_path + mqtt_config.get('CertFile')
    mqtt_key_file = config_path + mqtt_config.get('KeyFile')

    print("Connecting to local TOR controller on port %s" % tor_controller_port)

    with Controller.from_port(port=tor_controller_port) as controller:

        protocol_info = stem.connection.get_protocolinfo(controller)

        stem.connection.authenticate_safecookie(controller, protocol_info.cookie_path)

        print("Creating TOR Hidden Service...")

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

        onion_address = "%s.onion" % service.service_id

        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"

    if mqtt_broker_using_tor:
        client.proxy_set(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)

    print("Connecting 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,
        '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)