/*
 * Copyright (c) 2012-2019 Isode Limited.
 * All rights reserved.
 * See the COPYING file for more information.
 */

/*
 * Copyright (c) 2011 Tobias Markmann
 * Licensed under the simplified BSD license.
 * See Documentation/Licenses/BSD-simplified.txt for more information.
 */

#include <Swiften/FileTransfer/SOCKS5BytestreamServerManager.h>

#include <memory>

#include <boost/bind.hpp>

#include <Swiften/Base/Log.h>
#include <Swiften/FileTransfer/SOCKS5BytestreamServer.h>
#include <Swiften/FileTransfer/SOCKS5BytestreamServerPortForwardingUser.h>
#include <Swiften/FileTransfer/SOCKS5BytestreamServerResourceUser.h>
#include <Swiften/Network/ConnectionServer.h>
#include <Swiften/Network/ConnectionServerFactory.h>
#include <Swiften/Network/NATTraversalForwardPortRequest.h>
#include <Swiften/Network/NATTraversalGetPublicIPRequest.h>
#include <Swiften/Network/NATTraversalRemovePortForwardingRequest.h>
#include <Swiften/Network/NATTraverser.h>
#include <Swiften/Network/NetworkEnvironment.h>

using namespace Swift;

static const int LISTEN_PORTS_BEGIN = 10000;
static const int LISTEN_PORTS_END = 11000;

SOCKS5BytestreamServerManager::SOCKS5BytestreamServerManager(
        SOCKS5BytestreamRegistry* bytestreamRegistry,
        ConnectionServerFactory* connectionServerFactory,
        NetworkEnvironment* networkEnvironment,
        NATTraverser* natTraverser) :
            bytestreamRegistry(bytestreamRegistry),
            connectionServerFactory(connectionServerFactory),
            networkEnvironment(networkEnvironment),
            natTraverser(natTraverser),
            state(Start),
            server(nullptr),
            attemptedPortMapping_(false) {
}

SOCKS5BytestreamServerManager::~SOCKS5BytestreamServerManager() {
    SWIFT_LOG_ASSERT(!connectionServer, warning);
    SWIFT_LOG_ASSERT(!getPublicIPRequest, warning);
    SWIFT_LOG_ASSERT(!forwardPortRequest, warning);
    SWIFT_LOG_ASSERT(state == Start, warning);
    if (portMapping && !unforwardPortRequest) {
        SWIFT_LOG(warning) << "Port forwarding still alive. Trying to remove it now.";
        unforwardPortRequest = natTraverser->createRemovePortForwardingRequest(portMapping.get().getLocalPort(), portMapping.get().getPublicPort());
        unforwardPortRequest->start();
    }
}

std::shared_ptr<SOCKS5BytestreamServerResourceUser> SOCKS5BytestreamServerManager::aquireResourceUser() {
    std::shared_ptr<SOCKS5BytestreamServerResourceUser> resourceUser;
    if (s5bServerResourceUser_.expired()) {
        resourceUser = std::make_shared<SOCKS5BytestreamServerResourceUser>(this);
        s5bServerResourceUser_ = resourceUser;
    }
    else {
        resourceUser = s5bServerResourceUser_.lock();
    }
    return resourceUser;
}

std::shared_ptr<SOCKS5BytestreamServerPortForwardingUser> SOCKS5BytestreamServerManager::aquirePortForwardingUser() {
    std::shared_ptr<SOCKS5BytestreamServerPortForwardingUser> portForwardingUser;
    if (s5bServerPortForwardingUser_.expired()) {
        portForwardingUser = std::make_shared<SOCKS5BytestreamServerPortForwardingUser>(this);
        s5bServerPortForwardingUser_ = portForwardingUser;
    }
    else {
        portForwardingUser = s5bServerPortForwardingUser_.lock();
    }
    return portForwardingUser;
}

std::vector<HostAddressPort> SOCKS5BytestreamServerManager::getHostAddressPorts() const {
    std::vector<HostAddressPort> result;
    if (connectionServer) {
        std::vector<NetworkInterface> networkInterfaces = networkEnvironment->getNetworkInterfaces();
        for (const auto& networkInterface : networkInterfaces) {
            for (const auto& address : networkInterface.getAddresses()) {
                result.push_back(HostAddressPort(address, connectionServerPort));
            }
        }
    }
    return result;
}

std::vector<HostAddressPort> SOCKS5BytestreamServerManager::getAssistedHostAddressPorts() const {
    std::vector<HostAddressPort> result;
    if (publicAddress && portMapping) {
        result.push_back(HostAddressPort(*publicAddress, portMapping->getPublicPort()));
    }
    return result;
}

bool SOCKS5BytestreamServerManager::isInitialized() const {
    return state == Initialized;
}

void SOCKS5BytestreamServerManager::initialize() {
    if (state == Start) {
        state = Initializing;

        // Find a port to listen on
        assert(!connectionServer);
        unsigned short port;
        for (port = LISTEN_PORTS_BEGIN; port < LISTEN_PORTS_END; ++port) {
            SWIFT_LOG(debug) << "Trying to start server on port " << port;
            connectionServer = connectionServerFactory->createConnectionServer(HostAddress::fromString("::").get(), port);
            boost::optional<ConnectionServer::Error> error = connectionServer->tryStart();
            if (!error) {
                break;
            }
            else if (*error != ConnectionServer::Conflict) {
                SWIFT_LOG(debug) << "Error starting server";
                onInitialized(false);
                return;
            }
            connectionServer.reset();
        }
        if (!connectionServer) {
            SWIFT_LOG(debug) << "Unable to find an open port";
            onInitialized(false);
            return;
        }
        SWIFT_LOG(debug) << "Server started succesfully";
        connectionServerPort = port;

        // Start bytestream server. Should actually happen before the connectionserver is started
        // but that doesn't really matter here.
        assert(!server);
        server = new SOCKS5BytestreamServer(connectionServer, bytestreamRegistry);
        server->start();
        checkInitializeFinished();
    }
}

bool SOCKS5BytestreamServerManager::isPortForwardingReady() const {
    return attemptedPortMapping_ && !getPublicIPRequest && !forwardPortRequest;
}

void SOCKS5BytestreamServerManager::setupPortForwarding() {
    assert(server);
    attemptedPortMapping_ = true;

    // Retrieve public addresses
    assert(!getPublicIPRequest);
    publicAddress = boost::optional<HostAddress>();
    if ((getPublicIPRequest = natTraverser->createGetPublicIPRequest())) {
        getPublicIPRequest->onResult.connect(
                boost::bind(&SOCKS5BytestreamServerManager::handleGetPublicIPResult, this, _1));
        getPublicIPRequest->start();
    }

    // Forward ports
    auto port = server->getAddressPort().getPort();
    assert(!forwardPortRequest);
    portMapping = boost::optional<NATPortMapping>();
    if ((forwardPortRequest = natTraverser->createForwardPortRequest(port, port))) {
        forwardPortRequest->onResult.connect(
                boost::bind(&SOCKS5BytestreamServerManager::handleForwardPortResult, this, _1));
        forwardPortRequest->start();
    }
}

void SOCKS5BytestreamServerManager::removePortForwarding() {
    // remove port forwards
    if (portMapping) {
        unforwardPortRequest = natTraverser->createRemovePortForwardingRequest(portMapping.get().getLocalPort(), portMapping.get().getPublicPort());
        unforwardPortRequest->onResult.connect(boost::bind(&SOCKS5BytestreamServerManager::handleUnforwardPortResult, this, _1));
        unforwardPortRequest->start();
    }
}

void SOCKS5BytestreamServerManager::stop() {
    if (getPublicIPRequest) {
        getPublicIPRequest->stop();
        getPublicIPRequest.reset();
    }
    if (forwardPortRequest) {
        forwardPortRequest->stop();
        forwardPortRequest.reset();
    }
    if (unforwardPortRequest) {
        unforwardPortRequest->stop();
        unforwardPortRequest.reset();
    }
    if (server) {
        server->stop();
        delete server;
        server = nullptr;
    }
    if (connectionServer) {
        connectionServer->stop();
        connectionServer.reset();
    }

    state = Start;
}

void SOCKS5BytestreamServerManager::handleGetPublicIPResult(boost::optional<HostAddress> address) {
    if (address) {
        SWIFT_LOG(debug) << "Public IP discovered as " << address.get().toString() << ".";
    }
    else {
        SWIFT_LOG(debug) << "No public IP discoverable.";
    }

    publicAddress = address;

    getPublicIPRequest->stop();
    getPublicIPRequest.reset();
}

void SOCKS5BytestreamServerManager::handleForwardPortResult(boost::optional<NATPortMapping> mapping) {
    if (mapping) {
        SWIFT_LOG(debug) << "Mapping port was successful.";
    }
    else {
        SWIFT_LOG(debug) << "Mapping port has failed.";
    }

    portMapping = mapping;
    onPortForwardingSetup(mapping.is_initialized());

    forwardPortRequest->stop();
    forwardPortRequest.reset();
}

void SOCKS5BytestreamServerManager::handleUnforwardPortResult(boost::optional<bool> result) {
    if (result.is_initialized() && result.get()) {
        portMapping.reset();
    }
    else {
        SWIFT_LOG(warning) << "Failed to remove port forwarding.";
    }
    attemptedPortMapping_ = false;
    unforwardPortRequest.reset();
}

void SOCKS5BytestreamServerManager::checkInitializeFinished() {
    assert(state == Initializing);
    state = Initialized;
    onInitialized(true);
}