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 = mqtt_config.get('CaFile', fallback=None)
    mqtt_ca_file = config_path + mqtt_ca_file
    mqtt_cert_file = mqtt_config.get('CertFile', fallback=None)
    mqtt_cert_file = config_path + mqtt_cert_file
    mqtt_key_file = mqtt_config.get('KeyFile', fallback=None)
    mqtt_key_file = config_path + mqtt_key_file

    mqtt_use_tls = \
        mqtt_ca_file is not None and \
        mqtt_cert_file is not None and \
        mqtt_key_file is not None

    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

    cert_required = ssl.CERT_OPTIONAL
    if mqtt_require_certificate:
        cert_required = ssl.CERT_REQUIRED

    if mqtt_broker_using_tor:
        cert_required = ssl.CERT_OPTIONAL
        proxy_args = {
            'proxy_type': socks.SOCKS5,
            'proxy_addr': tor_proxy_host,
            'proxy_port': tor_proxy_port
        }

    if mqtt_use_tls:
        protocol = "mqtts"
        tls_args = {
            'ca_certs': mqtt_ca_file,
            'certfile': mqtt_cert_file,
            'keyfile': mqtt_key_file,
            'cert_reqs': 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)