/* -*- c-basic-offset: 8 -*-
   rdesktop: A Remote Desktop Protocol client.
   Master/Slave remote controlling
   Copyright 2013-2017 Henrik Andersson <hean01@cendio.se> for Cendio AB

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 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 General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#include "rdesktop.h"
#include "ssl.h"

#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <errno.h>
#include <limits.h>
#include <unistd.h>

#define CTRL_LINEBUF_SIZE 1024
#define CTRL_RESULT_SIZE 32
#define RDESKTOP_CTRLSOCK_STORE "/.local/share/rdesktop/ctrl"

#define CTRL_HASH_FLAG_SEAMLESS  1

#define ERR_RESULT_OK 0x00
#define ERR_RESULT_NO_SUCH_COMMAND 0xffffffff

extern RD_BOOL g_seamless_rdp;
extern uint8 g_static_rdesktop_salt_16[];
extern char g_codepage[16];

static RD_BOOL _ctrl_is_slave;
static int ctrlsock;
static char ctrlsock_name[PATH_MAX];
static struct _ctrl_slave_t *_ctrl_slaves;

#define CMD_SEAMLESS_SPAWN "seamless.spawn"

typedef struct _ctrl_slave_t
{
    struct _ctrl_slave_t *prev, *next;
    int sock;
    char linebuf[CTRL_LINEBUF_SIZE];
} _ctrl_slave_t;

static void _ctrl_slave_new(int sock);
static void _ctrl_slave_disconnect(int sock);
static void _ctrl_command_result(_ctrl_slave_t *slave, int result);
static void _ctrl_dispatch_command(_ctrl_slave_t *slave);
static RD_BOOL _ctrl_verify_unix_socket(void);
static void _ctrl_create_hash(const char *user, const char *domain, const char *host, char *hash, size_t hsize);

static void _ctrl_slave_new(int sock)
{
    _ctrl_slave_t *it, *ns;

    ns = (_ctrl_slave_t *)malloc(sizeof(_ctrl_slave_t));
    if (!ns)
    {
        logger(Core, Error, "_ctrl_slave_new(), malloc() failed");
        return;
    }
    memset(ns, 0, sizeof(_ctrl_slave_t));
    ns->sock = sock;

    it = _ctrl_slaves;

    while (it && it->next)
        it = it->next;

    if (it)
    {
        it->next = ns;
        ns->prev = it;
    }
    else
    {
        _ctrl_slaves = ns;
    }
}

static void _ctrl_slave_disconnect(int sock)
{
    _ctrl_slave_t *it;

    if (!_ctrl_slaves)
        return;

    it = _ctrl_slaves;

    while (it->next && it->sock != sock)
        it = it->next;

    if (it->sock == sock)
    {
        shutdown(sock, SHUT_RDWR);
        close(sock);

        if (it == _ctrl_slaves)
        {
            if (it->next)
                _ctrl_slaves = it->next;
            else
                _ctrl_slaves = NULL;
        }

        if (it->prev)
        {
            (it->prev)->next = it->next;
            if (it->next)
                (it->next)->prev = it->prev;
        }
        else if (it->next)
            (it->next)->prev = NULL;

        free(it);
    }
}

static void _ctrl_command_result(_ctrl_slave_t *slave, int result)
{
    char buf[64] = {0};

    if (result == 0)
        send(slave->sock, "OK\n", 3, 0);
    else
    {
        snprintf(buf, sizeof(buf), "ERROR %x\n", result);
        send(slave->sock, buf, strlen(buf), 0);
    }
}

static void _ctrl_dispatch_command(_ctrl_slave_t *slave)
{
    char *cmd, *p;
    unsigned int res;

    cmd = utils_string_unescape(slave->linebuf);
    if (!cmd)
    {
        _ctrl_command_result(slave, ERR_RESULT_NO_SUCH_COMMAND);
        return;
    }

    if (strncmp(cmd, CMD_SEAMLESS_SPAWN " ", strlen(CMD_SEAMLESS_SPAWN) + 1) == 0)
    {
        p = strstr(cmd, "seamlessrdpshell.exe");
        if (p)
            p += strlen("seamlessrdpshell.exe") + 1;
        else
            p = cmd + strlen(CMD_SEAMLESS_SPAWN) + 1;

        res = ERR_RESULT_OK;

        if (seamless_send_spawn(p) == (unsigned int)-1)
            res = 1;
    }
    else
    {
        res = ERR_RESULT_NO_SUCH_COMMAND;
    }
    free(cmd);

    _ctrl_command_result(slave, res);
}

static RD_BOOL _ctrl_verify_unix_socket(void)
{
    int s;
    struct sockaddr_un saun;
    socklen_t len;

    memset(&saun, 0, sizeof(struct sockaddr_un));

    if ((s = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
    {
        logger(Core, Error, "_ctrl_verify_unix_socket(), socket() failed: %s",
               strerror(errno));
        exit(1);
    }

    saun.sun_family = AF_UNIX;
    strncpy(saun.sun_path, ctrlsock_name, sizeof(saun.sun_path) - 1);
    len = sizeof(saun.sun_family) + strlen(saun.sun_path);

    if (connect(s, (struct sockaddr *)&saun, len) != 0)
    {
        close(s);
        return False;
    }

    shutdown(s, SHUT_RDWR);
    close(s);
    return True;
}

static void _ctrl_create_hash(const char *user, const char *domain, const char *host, char *hash, size_t hsize)
{
    RDSSL_SHA1 sha1;
    uint8 out[20], delim;
    uint16 version;
    uint32 flags;

    flags = 0;
    delim = '\0';
    version = 0x0100;

    if (g_seamless_rdp)
        flags = CTRL_HASH_FLAG_SEAMLESS;

    rdssl_sha1_init(&sha1);
    rdssl_sha1_update(&sha1, (uint8 *)&version, sizeof(version));
    rdssl_sha1_update(&sha1, &delim, 1);

    if (user)
        rdssl_sha1_update(&sha1, (uint8 *)user, strlen(user));
    rdssl_sha1_update(&sha1, &delim, 1);

    if (domain)
        rdssl_sha1_update(&sha1, (uint8 *)domain, strlen(domain));
    rdssl_sha1_update(&sha1, &delim, 1);

    if (host)
        rdssl_sha1_update(&sha1, (uint8 *)host, strlen(host));
    rdssl_sha1_update(&sha1, &delim, 1);

    rdssl_sha1_update(&sha1, (uint8 *)&flags, sizeof(flags));
    rdssl_sha1_final(&sha1, out);

    sec_hash_to_string(hash, hsize, out, sizeof(out));
}

int ctrl_init(const char *user, const char *domain, const char *host)
{
    struct stat st;
    struct sockaddr_un saun;
    char hash[41], path[PATH_MAX];
    const char *home;

    if (ctrlsock != 0 || _ctrl_is_slave)
        return 0;

    home = getenv("HOME");
    if (home == NULL)
    {
        return -1;
    }

    _ctrl_create_hash(user, domain, host, hash, sizeof(hash));
    snprintf(ctrlsock_name, sizeof(ctrlsock_name), "%s" RDESKTOP_CTRLSOCK_STORE "/%s.ctl", home, hash);

    snprintf(path, sizeof(path), "%s" RDESKTOP_CTRLSOCK_STORE, home);
    if (utils_mkdir_p(path, 0700) == -1)
    {
        logger(Core, Error, "ctrl_init(), utils_mkdir_p() failed: %s", strerror(errno));
        return -1;
    }

    if (stat(ctrlsock_name, &st) == 0)
    {
        if (_ctrl_verify_unix_socket() == True)
        {
            _ctrl_is_slave = True;
            return 1;
        }
        else
        {
            unlink(ctrlsock_name);
        }
    }

    if ((ctrlsock = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
    {
        logger(Core, Error, "ctrl_init(), socket() failed: %s", strerror(errno));
        exit(1);
    }

    memset(&saun, 0, sizeof(struct sockaddr_un));
    saun.sun_family = AF_UNIX;
    strncpy(saun.sun_path, ctrlsock_name, sizeof(saun.sun_path) - 1);
    if (bind(ctrlsock, (struct sockaddr *)&saun, sizeof(struct sockaddr_un)) < 0)
    {
        logger(Core, Error, "ctrl_init(), bind() failed: %s", strerror(errno));
        exit(1);
    }

    if (listen(ctrlsock, 5) < 0)
    {
        logger(Core, Error, "ctrl_init(), listen() failed: %s", strerror(errno));
        exit(1);
    }

    atexit(ctrl_cleanup);

    return 0;
}

void ctrl_cleanup(void)
{
    if (ctrlsock)
    {
        close(ctrlsock);
        unlink(ctrlsock_name);
    }
}

RD_BOOL ctrl_is_slave(void)
{
    return _ctrl_is_slave;
}

void ctrl_add_fds(int *n, fd_set *rfds)
{
    _ctrl_slave_t *it;
    if (ctrlsock == 0)
        return;

    FD_SET(ctrlsock, rfds);
    *n = MAX(*n, ctrlsock);

    it = _ctrl_slaves;
    while (it)
    {
        FD_SET(it->sock, rfds);
        *n = MAX(*n, it->sock);
        it = it->next;
    }
}

void ctrl_check_fds(fd_set *rfds, fd_set *wfds)
{
    UNUSED(wfds);
    int ns, res, offs;
    struct sockaddr_un fsaun;
    socklen_t fromlen;
    _ctrl_slave_t *it;

    if (ctrlsock == 0)
        return;

    memset(&fsaun, 0, sizeof(struct sockaddr_un));

    if (FD_ISSET(ctrlsock, rfds))
    {
        FD_CLR(ctrlsock, rfds);
        fromlen = sizeof(fsaun);
        ns = accept(ctrlsock, (struct sockaddr *)&fsaun, &fromlen);
        if (ns < 0)
        {
            logger(Core, Error, "ctrl_check_fds(), accept() failed: %s",
                   strerror(errno));
            exit(1);
        }

        _ctrl_slave_new(ns);
        return;
    }

    it = _ctrl_slaves;
    while (it)
    {
        if (FD_ISSET(it->sock, rfds))
        {
            offs = strlen(it->linebuf);
            res = recv(it->sock, it->linebuf + offs, CTRL_LINEBUF_SIZE - offs - 1, 0);
            FD_CLR(it->sock, rfds);

            if (it->linebuf[CTRL_LINEBUF_SIZE - 1] != '\0' &&
                it->linebuf[CTRL_LINEBUF_SIZE - 1] != '\n')
            {
                _ctrl_slave_disconnect(it->sock);
                break;
            }

            if (res > 0)
            {
                char *p;
                if ((p = strchr(it->linebuf, '\n')) == NULL)
                    continue;

                while (p)
                {
                    if (p > it->linebuf && *(p - 1) != '\\')
                        break;
                    p = strchr(p + 1, '\n');
                }

                if (p == NULL)
                    continue;

                *p = '\0';
                _ctrl_dispatch_command(it);
                memset(it->linebuf, 0, CTRL_LINEBUF_SIZE);
            }
            else
            {
                _ctrl_slave_disconnect(it->sock);
                break;
            }
        }
        it = it->next;
    }
}

int ctrl_send_command(const char *cmd, const char *arg)
{
    FILE *fp;
    struct sockaddr_un saun;
    int s, len, index, ret;
    char data[CTRL_LINEBUF_SIZE], tmp[CTRL_LINEBUF_SIZE];
    char result[CTRL_RESULT_SIZE], c, *escaped;

    escaped = NULL;

    if (!_ctrl_is_slave)
        return -1;

    if ((s = socket(AF_UNIX, SOCK_STREAM, 0)) < 0)
    {
        logger(Core, Error, "ctrl_send_command(), socket() failed: %s", strerror(errno));
        exit(1);
    }

    memset(&saun, 0, sizeof(struct sockaddr_un));
    saun.sun_family = AF_UNIX;
    strncpy(saun.sun_path, ctrlsock_name, sizeof(saun.sun_path) - 1);
    len = sizeof(saun.sun_family) + strlen(saun.sun_path);

    if (connect(s, (struct sockaddr *)&saun, len) < 0)
    {
        logger(Core, Error, "ctrl_send_command(), connect() failed: %s", strerror(errno));
        exit(1);
    }

    snprintf(data, sizeof(data), "%s %s", cmd, arg);
    ret = utils_locale_to_utf8(data, strlen(data), tmp, sizeof(tmp) - 1);

    if (ret != 0)
        goto bail_out;

    escaped = utils_string_escape(tmp);
    if ((strlen(escaped) + 1) > sizeof(data) - 1)
        goto bail_out;

    send(s, escaped, strlen(escaped), 0);
    send(s, "\n", 1, 0);

    fp = fdopen(s, "r");
    index = 0;
    while ((c = fgetc(fp)) != EOF && index < CTRL_RESULT_SIZE && c != '\n')
    {
        result[index] = c;
        index++;
    }
    result[index - 1] = '\0';

    if (strncmp(result, "ERROR ", 6) == 0)
    {
        if (sscanf(result, "ERROR %d", &ret) != 1)
            ret = -1;
    }

bail_out:
    free(escaped);
    shutdown(s, SHUT_RDWR);
    close(s);

    return ret;
}