/***************************************************************************
 *   Copyright 2011 David Nolden <david.nolden.kdevelop@art-master.de>                         *
 *                                                                         *
 *   This program is free software; you can redistribute it and/or modify  *
 *   it under the terms of the GNU Library General Public License as       *
 *   published by the Free Software Foundation; either version 2 of the    *
 *   License, or (at your option) any later version.                       *
 *                                                                         *
 *   This program is distributed in the hope that it will be useful,       *
 *   but WITHOUT ANY WARRANTY; without even the implied warranty of        *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         *
 *   GNU General Public License for more details.                          *
 *                                                                         *
 *   You should have received a copy of the GNU Library General Public     *
 *   License along with this program; if not, write to the                 *
 *   Free Software Foundation, Inc.,                                       *
 *   51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.         *
 ***************************************************************************/

#include <cassert>
#include <cerrno>
#include <cstdlib>
#include <cstring>
#include <fcntl.h>
#include <iostream>
#include <netdb.h>
#include <netinet/in.h>
#include <sstream>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <unistd.h>

#ifndef HAVE_MSG_NOSIGNAL
#define MSG_NOSIGNAL 0
#endif

/**
 * The goal of this utility is transforming the abstract unix-socket which is used by dbus
 * into a TCP socket which can be forwarded to a target machine by ssh tunneling, and then on
 * the target machine back into an abstract unix socket.
 *
 * This tool basically works similar to the "socat" utility, except that it works properly
 * for this special case. It is merely responsible for the transformation between abstract unix
 * sockets and tcp sockets.
 *
 * Furthermore, this tool makes the 'EXTERNAL' dbus authentication mechanism work even across
 * machines with different user IDs.
 *
 * This is how the EXTERNAL mechanism works (I found this in a comment of some ruby dbus library):
 *   Take the user id (eg integer 1000) make a string out of it "1000", take
 *   each character and determine hex value "1" => 0x31, "0" => 0x30. You
 *   obtain for "1000" => 31303030 This is what the server is expecting.
 *   Why? I dunno. How did I come to that conclusion? by looking at rbus
 *   code. I have no idea how he found that out.
 *
 * The dbus client performs the EXTERNAL authentication by sending "AUTH EXTERNAL 31303030\r\n" once
 * after opening the connection, so we can "repair" the authentication by overwriting the token in that
 * string through the correct one.
 * */

const bool debug = false;

/**
 * Returns the valid dbus EXTERNAL authentication token for the current user (see above)
 * */
std::string getAuthToken()
{
    // Get uid
    int uid = getuid();

    std::ostringstream uidStream;
    uidStream << uid;

    std::string uidStr = uidStream.str();

    const char hexdigits[16] = {
        '0', '1', '2', '3', '4', '5', '6', '7', '8', '9',
        'a', 'b', 'c', 'd', 'e', 'f'
    };

    std::ostringstream hexStream;
    for (char c : uidStr) {
        auto byte = (unsigned char)c;
        hexStream << hexdigits[byte >> 4] << hexdigits[byte & 0x0f];
    }

    return hexStream.str();
}

/**
 * Shuffles all data between the two file-descriptors until one of them fails or reaches EOF.
 * */
void shuffleBetweenStreams(int side1, int side2, bool fixSide1AuthToken)
{
    char buffer[1000];
    char buffer2[1000];

    // Set non-blocking mode
    int opts = fcntl(side1, F_GETFL);
    opts |= O_NONBLOCK;
    fcntl(side1, F_SETFL, opts);

    opts = fcntl(side2, F_GETFL);
    opts |= O_NONBLOCK;
    fcntl(side2, F_SETFL, opts);

    while (true) {
        int r1 = read(side1, buffer, 500); // We read less than 1000, so we have same additional space when changing the auth token
        int r2 = read(side2, buffer2, 500);

        if (r1 < -1 || r1 == 0) {
            if (debug)
                std::cerr << "stream 1 failed: " << r1 << std::endl;
            return;
        }
        if (r2 < -1 || r2 == 0) {
            if (debug)
                std::cerr << "stream 2 failed: " << r2 << std::endl;
            return;
        }

        if (r1 > 0) {
            if (debug)
                std::cerr << "transferring " << r1 << " from 1 to 2" << std::endl;

            if (fixSide1AuthToken) {
                if (r1 > 15 && memcmp(buffer, "\0AUTH EXTERNAL ", 15) == 0) {
                    int endPos = -1;
                    for (int i = 15; i < r1; ++i) {
                        if (buffer[i] == '\r') {
                            endPos = i;
                            break;
                        }
                    }

                    if (endPos != -1) {
                        std::string oldToken = std::string(buffer + 15, endPos - 15);
                        std::string newToken = getAuthToken();

                        int difference = newToken.size() - oldToken.size();
                        r1 += difference;
                        assert(r1 > 0 && r1 <= 1000);
                        memmove(buffer + endPos + difference, buffer + endPos, r1 - difference - endPos);
                        memcpy(buffer + 15, newToken.data(), newToken.size());
                        assert(buffer[endPos + difference] == '\r');
                        assert(buffer[endPos + difference - 1] == newToken[newToken.size() - 1]);
                    } else {
                        std::cout << "could not fix auth token, not enough data available" << std::endl;
                    }
                } else {
                    std::cout << "could not fix auth token" << std::endl;
                }
                fixSide1AuthToken = false;
            }

            opts = fcntl(side2, F_GETFL);
            opts ^= O_NONBLOCK;
            fcntl(side2, F_SETFL, opts);

            int w2 = send(side2, buffer, r1, MSG_NOSIGNAL);

            if (w2 < 0) {
                if (debug)
                    std::cerr << "writing to side 2 failed, ending: " << w2 << std::endl;
                return;
            }
            assert(w2 == r1);

            opts = fcntl(side2, F_GETFL);
            opts |= O_NONBLOCK;
            fcntl(side2, F_SETFL, opts);
        }

        if (r2 > 0) {
            if (debug)
                std::cerr << "transferring " << r2 << " from 2 to 1" << std::endl;
            opts = fcntl(side1, F_GETFL);
            opts ^= O_NONBLOCK;
            fcntl(side1, F_SETFL, opts);

            int w1 = send(side1, buffer2, r2, MSG_NOSIGNAL);

            if (w1 < 0) {
                if (debug)
                    std::cerr << "writing to side 1 failed, ending: " << w1 << std::endl;
                return;
            }
            assert(w1 == r2);

            opts = fcntl(side1, F_GETFL);
            opts |= O_NONBLOCK;
            fcntl(side1, F_SETFL, opts);
        }
        usleep(1000);
    }
}

int main(int argc, char** argv)
{
    int serverfd;

    if (argc < 2) {
        std::cerr << "need arguments:" << std::endl;
        std::cerr <<
        "[port]                    -   Open a server on this TCP port and forward connections to the local DBUS session"
            "                              (the DBUS_SESSION_BUS_ADDRESS environment variable must be set)";
        std::cerr <<
        "[port]  [fake dbus path]  -   Open a server on the fake dbus path and forward connections to the given local TCP port";
        std::cerr << ""
            "The last argument may be the --bind-only option, in which case the application only tries to"
            "open the server, but does not wait for clients to connect. This is useful to test whether the"
            "server port/path is available.";
        return 10;
    }

    bool waitForClients = true;

    if (std::string(argv[argc - 1]) == "--bind-only") {
        waitForClients = false;
        argc -= 1;
    }

    std::string dbusAddress(getenv("DBUS_SESSION_BUS_ADDRESS"));

    std::string path;

    if (argc == 2) {
        if (waitForClients && debug)
            std::cout << "forwarding from the local TCP port " << argv[1] << " to the local DBUS session at " <<
            dbusAddress.data() << std::endl;

        if (dbusAddress.empty()) {
            std::cerr << "The DBUS_SESSION_BUS_ADDRESS environment variable is not set" << std::endl;
            return 1;
        }

        // Open a TCP server

        std::string abstractPrefix("unix:abstract=");

        if (dbusAddress.substr(0, abstractPrefix.size()) != abstractPrefix) {
            std::cerr << "DBUS_SESSION_BUS_ADDRESS does not seem to use an abstract unix domain socket as expected" <<
            std::endl;
            return 2;
        }

        path = dbusAddress.substr(abstractPrefix.size(), dbusAddress.size() - abstractPrefix.size());
        if (path.find(",guid=") != std::string::npos)
            path = path.substr(0, path.find(",guid="));

        // Mark it as an abstract unix domain socket
        serverfd = socket(AF_INET, SOCK_STREAM, 0);

        if (serverfd < 0) {
            if (waitForClients)
                std::cerr << "ERROR opening server socket" << std::endl;
            return 3;
        }

        int portno = atoi(argv[1]);

        sockaddr_in server_addr;

        memset(&server_addr, 0, sizeof(server_addr));
        server_addr.sin_family = AF_INET;
        server_addr.sin_addr.s_addr = INADDR_ANY;
        server_addr.sin_port = htons(portno);

        if (bind(serverfd, ( struct sockaddr* ) &server_addr,
                 sizeof(server_addr)) < 0) {
            if (waitForClients)
                std::cerr << "ERROR opening the server" << std::endl;
            return 7;
        }
    } else if (argc == 3) {
        if (waitForClients && debug)
            std::cout << "forwarding from the local abstract unix domain socket " << argv[2] <<
            " to the local TCP port " << argv[1] << std::endl;
        // Open a unix domain socket server
        serverfd = socket(AF_UNIX, SOCK_STREAM, 0);

        if (serverfd < 0) {
            if (waitForClients)
                std::cerr << "ERROR opening server socket" << std::endl;
            return 3;
        }

        path = std::string(argv[2]);

        sockaddr_un serv_addr;
        memset(&serv_addr, 0, sizeof(serv_addr));
        serv_addr.sun_family = AF_UNIX;
        serv_addr.sun_path[0] = '\0'; // Mark as an abstract socket
        strcpy(serv_addr.sun_path + 1, path.data());

        if (debug)
            std::cout << "opening at " << path.data() << std::endl;

        if (bind(serverfd, ( sockaddr* ) &serv_addr, sizeof (serv_addr.sun_family) + 1 + path.length()) < 0) {
            if (waitForClients)
                std::cerr << "ERROR opening the server" << std::endl;
            return 7;
        }
    } else {
        std::cerr << "Wrong arguments";
        return 1;
    }

    listen(serverfd, 10);

    while (waitForClients) {
        if (debug)
            std::cerr << "waiting for client" << std::endl;
        sockaddr_in cli_addr;
        socklen_t clilen = sizeof(cli_addr);
        int connectedclientsockfd = accept(serverfd,
                                           ( struct sockaddr* ) &cli_addr,
                                           &clilen);

        if (connectedclientsockfd < 0) {
            std::cerr << "ERROR on accept" << std::endl;
            return 8;
        }

        if (debug)
            std::cerr << "got client" << std::endl;

        int sockfd;

        int addrSize;
        sockaddr* useAddr = nullptr;
        sockaddr_un serv_addru;
        sockaddr_in serv_addrin;

        if (argc == 2) {
            sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
            if (sockfd < 0) {
                std::cerr << "ERROR opening socket" << std::endl;
                return 3;
            }
            memset(&serv_addru, 0, sizeof(serv_addru));
            serv_addru.sun_family = AF_UNIX;
            serv_addru.sun_path[0] = '\0'; // Mark as an abstract socket
            strcpy(serv_addru.sun_path + 1, path.data());
            addrSize = sizeof (serv_addru.sun_family) + 1 + path.size();
            useAddr = ( sockaddr* )&serv_addru;

            if (debug)
                std::cout << "connecting to " << path.data() << std::endl;
        } else {
            sockfd = socket(AF_INET, SOCK_STREAM, 0);
            if (sockfd < 0) {
                std::cerr << "ERROR opening socket" << std::endl;
                return 3;
            }
            int port = atoi(argv[1]);
            hostent* server = gethostbyname("localhost");
            if (server == nullptr) {
                std::cerr << "failed to get server" << std::endl;
                return 5;
            }
            memset(&serv_addrin, 0, sizeof(serv_addrin));
            serv_addrin.sin_family = AF_INET;
            serv_addrin.sin_addr.s_addr = INADDR_ANY;
            serv_addrin.sin_port = htons(port);
            memcpy(&serv_addrin.sin_addr.s_addr, server->h_addr, server->h_length);
            addrSize = sizeof (serv_addrin);
            useAddr = ( sockaddr* )&serv_addrin;

            if (debug)
                std::cout << "connecting to port " << port << std::endl;
        }

        if (connect(sockfd, useAddr, addrSize) < 0) {
            int res = errno;
            if (res == ECONNREFUSED)
                std::cerr << "ERROR while connecting: connection refused" << std::endl;
            else if (res == ENOENT)
                std::cerr << "ERROR while connecting: no such file or directory" << std::endl;
            else
                std::cerr << "ERROR while connecting" << std::endl;
            return 5;
        }

        shuffleBetweenStreams(connectedclientsockfd, sockfd, argc == 2);
        close(sockfd);
        close(connectedclientsockfd);
    }
    return 0;
}
