from stem.control import Controller
import stem.connection
import paho.mqtt.client as mqtt
from paho.mqtt.client 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()
    config.read(config_path + "torch.conf")

    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)
    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')

    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("Connected to Tor on port %s" % tor_controller_port)

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

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

        print("Created Tor Hidden Service for local port %s at %s" % (ssh_port, onion_address))

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

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

    if mqtt_require_certificate:
        client.tls_set(
            ca_certs=mqtt_ca_file,
            certfile=mqtt_cert_file,
            keyfile=mqtt_key_file,
            cert_reqs=ssl.CERT_REQUIRED)
        protocol = "mqtts"

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

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

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