/*
 * © Copyright 1996-2012 ECMWF.
 *
 * This software is licensed under the terms of the Apache Licence Version 2.0
 * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
 * In applying this licence, ECMWF does not waive the privileges and immunities
 * granted to it by virtue of its status as an intergovernmental organisation nor
 * does it submit to any jurisdiction.
 */

#include <ctype.h>
#include "mars.h"


static request* proxies = 0;
static int inited       = 0;


static char* _config(const request* r, const char* name) {
    const char* v = get_value(r, name, 0);
    if (v) {
        return strdup(v);
    }
    print_one_request(r);
    marslog(LOG_EXIT, "Missing '%s' in proxy configuration", name);
    return 0;
}

const request* find_proxy(const char* name) {
    if (!inited) {
        inited = 1;
        if (mars.proxies) {
            proxies = read_request_file(mars.proxies);
        }
    }


    request* r = proxies;
    while (r) {
        const char* n = _config(r, "name");

        if (n && strcmp(n, name) == 0) {
            return r;
        }
        r = r->next;
    }
    return 0;
}


static int http_proxy_connector(const char* host, int port, int retries, const request* config) {

    const char* proxy_host = _config(config, "host");
    int proxy_port         = atoi(_config(config, "port"));

    if (mars.show_hosts) {
        marslog(LOG_INFO, "Connecting to %s:%d using HTTP CONNECT proxy on %s:%d",
                host, port, proxy_host, proxy_port);
    }

    // https://github.com/jab/trio_http_proxy
    // https://en.wikipedia.org/wiki/HTTP_tunnel#HTTP_CONNECT_method
    int s = call_server(proxy_host,
                        proxy_port,
                        retries,
                        NULL);

    char buf[1024];
    sprintf(buf, "CONNECT %s:%d HTTP/1.0\r\n", host, port);
    if (get_value(config, "authorization", 0)) {
        strcat(buf, "Proxy-Authorization: ");
        strcat(buf, _config(config, "authorization"));
        strcat(buf, "\r\n");
    }
    strcat(buf, "\r\n");

    char* p = buf;

    if (s < 0) {
        marslog(LOG_EROR, "Cannot connect to proxy");
        return -1;
    }

    if (writetcp(&s, buf, strlen(buf)) != strlen(buf)) {
        marslog(LOG_EROR, "Cannot write to proxy");
        close(s);
        return -1;
    }

    int state = 0;
    int i     = 0;
    buf[0]    = 0;
    char c;
    int code = 0;

    for (;;) {
        if (readtcp(&s, &c, 1) != 1) {
            marslog(LOG_EROR, "Cannot read from proxy");
            close(s);
            return -1;
        }

        switch (c) {
            case '\n':
                if (state == 1) {
                    if (mars.show_hosts) {
                        marslog(LOG_DBUG, "Proxy response <%s>", buf);
                    }
                    i = 0;
                    if (code == 0) {
                        char* p = buf;
                        while (*p && !isspace(*p))
                            p++;
                        code = atoi(p);
                    }
                }
                if (state == 1 || state == 3) {
                    state++;
                }
                else {
                    state = 0;
                }
                break;

            case '\r':
                state++;
                break;

            default:
                state = 0;
                if (i < sizeof(buf) - 1) {
                    buf[i++] = c;
                    buf[i]   = 0;
                }
                else {
                    marslog(LOG_EROR, "Proxy response too long");
                    close(s);
                    return -1;
                }
                break;
        }


        if (state == 4) {
            break;
        }
    }

    if (code != 200) {
        marslog(LOG_EROR, "Proxy response code %d", code);
        close(s);
        return -1;
    }

    marslog(LOG_INFO, "Connected to proxy");

    return s;
}


static const char* sock_statuses[] = {
    "request granted",
    "general failure",
    "connection not allowed by ruleset",
    "network unreachable",
    "host unreachable",
    "connection refused by destination host",
    "TTL expired",
    "command not supported / protocol error",
    "address type not supported",
};

static int socks5_connector(const char* host, int port, int retries, const request* config) {
    // https://en.wikipedia.org/wiki/SOCKS
    // https://github.com/rofl0r/microsocks
    // https://linux.die.net/man/8/socks

    const char* proxy_host = _config(config, "host");
    int proxy_port         = atoi(_config(config, "port"));
    const char* username;
    const char* password;

    if (mars.show_hosts) {
        marslog(LOG_INFO, "Connecting to %s:%d using SOCKS proxy on %s:%d",
                host, port, proxy_host, proxy_port);
    }

    int s = call_server(proxy_host,
                        proxy_port,
                        retries,
                        NULL);


    char buf[1024];
    int i = 0;
    int len;
    const char* p;

    // Client greeting

    buf[i++] = 0x05;  // VER
    buf[i++] = 0x02;  // NMETHODS
    buf[i++] = 0x00;  // No authentication
    buf[i++] = 0x02;  // Username/password

    if (writetcp(&s, buf, i) != i) {
        marslog(LOG_EROR, "Cannot write to proxy (Client greeting)");
        close(s);
        return -1;
    }

    // Server greeting
    if (readtcp(&s, buf, 2) != 2) {
        marslog(LOG_EROR | LOG_PERR, "Cannot read 2 from proxy (Server greeting)");
        close(s);
        return -1;
    }

    if (buf[0] != 0x05) {  // VER
        marslog(LOG_EROR, "Proxy response code VER=%x", buf[0]);
        close(s);
        return -1;
    }

    switch (buf[1]) {
        case 0x00:  // No authentication
            break;

        case 0x02:  // Username/password

            username = _config(config, "username");
            password = _config(config, "password");

            i        = 0;
            buf[i++] = 0x01;  // VER
            buf[i++] = strlen(username);
            p        = username;
            while (*p) {
                buf[i++] = *p++;
            }
            buf[i++] = strlen(password);
            p        = password;
            while (*p) {
                buf[i++] = *p++;
            }
            if (writetcp(&s, buf, i) != i) {
                marslog(LOG_EROR, "Cannot write to proxy (Client authentication)");
                close(s);
                return -1;
            }
            if (readtcp(&s, buf, 2) != 2) {
                marslog(LOG_EROR | LOG_PERR, "Cannot read 2 from proxy (Server authentication)");
                close(s);
                return -1;
            }
            if (buf[0] != 0x01) {  // VER
                marslog(LOG_EROR, "Proxy response code VER=%x", buf[0]);
                close(s);
                return -1;
            }
            if (buf[1] != 0x00) {  // STATUS
                marslog(LOG_EROR, "Proxy response code STATUS=%x", buf[1]);
                close(s);
                return -1;
            }
            break;

        default:
            marslog(LOG_EROR, "Proxy response code CAUTH=%x", buf[1]);
            close(s);
            return -1;
    }
    i = 0;

    buf[i++] = 0x05;  // VER
    buf[i++] = 0x01;  // CMD establish a TCP/IP stream connection
    buf[i++] = 0x00;  // RSV
    buf[i++] = 0x03;  // ATYP  Domain name
    buf[i++] = strlen(host);

    p = host;
    while (*p) {
        buf[i++] = *p++;
    }
    buf[i++] = port >> 8;
    buf[i++] = port & 0xff;

    if (writetcp(&s, buf, i) != i) {
        marslog(LOG_EROR, "Cannot write to proxy (Client request)");
        close(s);
        return -1;
    }

    if (readtcp(&s, buf, 4) != 4) {
        marslog(LOG_EROR | LOG_PERR, "Cannot read 4 from proxy (Server response)");
        close(s);
        return -1;
    }

    if (buf[0] != 0x05) {  // VER
        marslog(LOG_EROR, "Proxy response code VER=%x", buf[0]);
        close(s);
        return -1;
    }

    if (buf[1] != 0x00) {  // STATUS
        const char* message = buf[1] < (sizeof(sock_statuses) / sizeof(sock_statuses[0])) ? sock_statuses[buf[1]] : "unknown";
        marslog(LOG_EROR, "Proxy response code STATUS=%x %s", buf[1], message);
        close(s);
        return -1;
    }

    if (buf[2] != 0x00) {  // RSV
        marslog(LOG_EROR, "Proxy response code RSV=%x", buf[1]);
        close(s);
        return -1;
    }

    switch (buf[3]) {
        case 0x01:  // IPv4
            if (readtcp(&s, buf, 4) != 4) {
                marslog(LOG_EROR, "Cannot read IPv4 from proxy");
                close(s);
                return -1;
            }
            break;

        case 0x03:  // Domain name
            if (readtcp(&s, buf, 1) != 1) {
                marslog(LOG_EROR, "Cannot read Domain name from proxy");
                close(s);
                return -1;
            }
            len = buf[0];
            if (readtcp(&s, buf, len) != len) {
                marslog(LOG_EROR, "Cannot read Domain name from proxy");
                close(s);
                return -1;
            }
            break;

        case 0x04:  // IPv6
            if (readtcp(&s, buf, 16) != 16) {
                marslog(LOG_EROR, "Cannot read IPv6 from proxy");
                close(s);
                return -1;
            }
            break;

        default:
            marslog(LOG_EROR, "Proxy response code for ADDR %x", buf[3]);
            close(s);
            return -1;
    }

    if (readtcp(&s, buf, 2) != 2) {  // PORT
        marslog(LOG_EROR, "Cannot read PORT from proxy");
        close(s);
        return -1;
    }

    marslog(LOG_DBUG, "Connected to SOCK proxy");

    return s;
}


tcp_connector proxy_for(const request* config) {

    const char* type = _config(config, "type");
    if (strcmp(type, "http") == 0) {
        return &http_proxy_connector;
    }
    if (strcmp(type, "socks") == 0) {
        return &socks5_connector;
    }

    marslog(LOG_EXIT, "Unknown proxy type %s", type);
    return 0;
}
