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

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

    protocol = "mqtt"
    tls_args = None
    proxy_args = None

    if mqtt_broker_using_tor:
        proxy_args = {
            'proxy_type': socks.SOCKS5,
            'proxy_addr': tor_proxy_host,
            'proxy_port': tor_proxy_port
        }
    else:
        if mqtt_require_certificate:
            protocol = "mqtts"
            tls_args = {
                'ca_certs': mqtt_ca_file,
                'certfile': mqtt_cert_file,
                'keyfile': mqtt_key_file,
                'cert_reqs': ssl.CERT_REQUIRED
            }

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

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

    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)