/*
  Copyright (C) 2013-2017 Grégory Soutadé
  
  This file is part of gPass.
  
  gPass 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.
  
  gPass 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 gPass.  If not, see <http://www.gnu.org/licenses/>.
*/

#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <signal.h>

#include <curl/curl.h>
#include <openssl/opensslv.h>
#include <openssl/evp.h>

#include "ini.h"

#define STRNCMP(a, b) strncmp(a, b, sizeof(b)-1)

#define DEFAULT_CONFIG_FILE ".local/share/gpass/gpass.ini"

#define DEFAULT_PBKDF2_LEVEL    1000
#define MASTER_KEY_LENGTH    (256/8)
#define GLOBAL_IV_LENGTH          16
#define BLOCK_SIZE           (128/8)
#define DEFAULT_SERVER_PORT      443
#define SERVER_PROTOCOL            4
#define RESPONSE_SIZE           2048
#define MAX_SUBDOMAINS            10
#define DISPLAY_TIME              30 // 30 seconds

struct gpass_parameters {
    unsigned      pbkdf2_level;
    char          *server;
    char          *salt;
    char          *domain;
    char          *username;
    char          *orig_master_key;
    unsigned char *derived_master_key;
    unsigned       server_port;
    unsigned       verbose;
    char          *ca_path;
    unsigned       verify_ssl_peer;
    unsigned       port_set;
    unsigned       crypto_v1_compatible;
    unsigned char *global_iv;
} ;

#if OPENSSL_VERSION_NUMBER >= 0x10010000
// OpenSSL >= 1.1
static EVP_MD_CTX * s_md_ctx;
#else
static EVP_MD_CTX * s_md_ctx;
static EVP_MD_CTX ss_md_ctx;
#define EVP_MD_CTX_new(...) &ss_md_ctx
#define EVP_MD_CTX_free(...)
#endif
static const EVP_MD * s_md_256;

static EVP_CIPHER_CTX * s_cipher_ctx;
static int s_stop_display = 0;

static void signal_handler(int signum)
{
    s_stop_display = 1;
}

static void display_password(char* password, int time)
{
    int print_len = 0;

    for (; time && !s_stop_display; time--)
    {
        print_len = printf("\r(%02d) Password found: %s", time, password);
        fflush(stdout);
        sleep(1);
    }

    // Clear line
    print_len++; // For C or Z
    printf("\r");
    while (print_len--)
        printf(" ");
    printf("\n");
}

static int digest(unsigned char** out, unsigned char* in, unsigned size)
{
    *out = NULL;
    EVP_DigestInit(s_md_ctx, s_md_256);
    EVP_DigestUpdate(s_md_ctx, in, size);
    *out = malloc(32);
    return EVP_DigestFinal(s_md_ctx, *out, NULL);
}

static void derive_master_key(struct gpass_parameters* params)
{
    if (!params->derived_master_key)
        params->derived_master_key = malloc(MASTER_KEY_LENGTH);

    if (!params->global_iv)
        params->global_iv = malloc(GLOBAL_IV_LENGTH);

    PKCS5_PBKDF2_HMAC(params->orig_master_key, strlen(params->orig_master_key),
                      (unsigned char*)params->salt, strlen(params->salt),
                      params->pbkdf2_level, EVP_sha256(),
                      MASTER_KEY_LENGTH, params->derived_master_key);

    PKCS5_PBKDF2_HMAC(params->salt, strlen(params->salt),
                      (unsigned char*)params->orig_master_key, strlen(params->orig_master_key),
                      params->pbkdf2_level, EVP_sha256(),
                      GLOBAL_IV_LENGTH, params->global_iv);
}

static void bin_to_hex(unsigned char* bin, unsigned char* hex, unsigned bin_size)
{
    unsigned char tmp;

    for (; bin_size--; bin++)
    {
        tmp = (*bin >> 4) & 0xf;
        if (tmp <= 9)
            *hex++ = '0' + tmp;
        else
            *hex++ = 'a' + (tmp-10);

        tmp = *bin  & 0xf;
        if (tmp <= 9)
            *hex++ = '0' + tmp;
        else
            *hex++ = 'a' + (tmp-10);
    }
}

static void hex_to_bin(unsigned char* bin, unsigned char* hex, long hex_size)
{
    unsigned char tmp;

    // Round to 2
    hex_size &= ~1;

    for (; hex_size; hex_size-=2, bin++)
    {
        tmp = *hex++;
        if (tmp >= '0' && tmp <= '9')
            *bin = (tmp - '0') << 4;
        else if (tmp >= 'a' && tmp <= 'f')
            *bin = ((tmp - 'a')+10) << 4;
        else
            *bin = ((tmp - 'A')+10) << 4;

        tmp = *hex++;
        if (tmp >= '0' && tmp <= '9')
            *bin |= (tmp - '0');
        else if (tmp >= 'a' && tmp <= 'f')
            *bin |= ((tmp - 'a')+10);
        else
            *bin |= ((tmp - 'A')+10);
    }
}

static void encrypt_domain_v1(struct gpass_parameters* params, char* domain,
                           unsigned char** res, unsigned* out_size)
{
    unsigned size = 2+strlen(domain)+1+strlen(params->username);
    unsigned char* buffer, *tmp;

    if (params->verbose)
        printf("%s: %s\n", __func__, domain);

    if ((size % BLOCK_SIZE))
        size = ((size/BLOCK_SIZE)+1)*BLOCK_SIZE;
    
    buffer = malloc(size+1); // Cause snprintf() add a final \0
    memset(buffer, 0, size+1);

    snprintf((char*)buffer, size+1, "@@%s;%s", domain, params->username);

    tmp = malloc(size);
    *res = malloc(size*2);

    EVP_EncryptInit(s_cipher_ctx, EVP_aes_256_ecb(), params->derived_master_key, NULL);
    EVP_CipherUpdate(s_cipher_ctx, tmp, (int*)out_size, buffer, size);

    bin_to_hex(tmp, *res, size);

    *out_size *= 2;

    free(buffer);
    free(tmp);
}

static void encrypt_domain(struct gpass_parameters* params, char* domain,
                           unsigned char** res, unsigned* out_size)
{
    unsigned size = strlen(domain)+1+strlen(params->username);
    unsigned padded_size;
    unsigned char* buffer, *tmp;

    if (params->verbose)
        printf("%s: %s\n", __func__, domain);

    if ((size % BLOCK_SIZE))
        size = ((size/BLOCK_SIZE)+1)*BLOCK_SIZE;
    padded_size = size;

    size += 16; // For digest
    
    buffer = malloc(size);
    memset(buffer, 0, size);

    snprintf((char*)buffer, size, "%s;%s", domain, params->username);

    // Append digest
    digest(&tmp, buffer, padded_size);
    memcpy(&buffer[padded_size], &tmp[8], 16);
    free(tmp);
    
    tmp = malloc(size);
    *res = malloc(size*2);

    EVP_EncryptInit(s_cipher_ctx, EVP_aes_256_cbc(), params->derived_master_key, params->global_iv);
    EVP_CipherUpdate(s_cipher_ctx, tmp, (int*)out_size, buffer, size);

    bin_to_hex(tmp, *res, size);

    *out_size *= 2;

    free(buffer);
    free(tmp);
}

static void append_to_request(char** request, char* new_req, unsigned new_req_size)
{
    static int cur_req_idx = 0;
    int size_added;

    if (!cur_req_idx)
    {
        *request = malloc(3+new_req_size+1);
        snprintf(*request, 3+new_req_size+1, "k0=%s", new_req);
    }
    else
    {
        size_added = 4+new_req_size;
        if (cur_req_idx >= 10)
            size_added++;

        *request = realloc(*request, strlen(*request)+1+size_added);

        snprintf(&((*request)[strlen(*request)]), size_added+1, "&k%d=%s",
                 cur_req_idx, new_req);
    }

    cur_req_idx++;
}

static char* wildcard_domain(char* domain)
{
    int cur_level = 1;
    char* level_ptr[MAX_SUBDOMAINS], *tmp, *res = NULL;
    int level_length[MAX_SUBDOMAINS];

    memset(level_ptr, 0, sizeof(level_ptr));
    memset(level_length, 0, sizeof(level_length));
    level_ptr[0] = domain;

    for (tmp=domain; *tmp && cur_level < MAX_SUBDOMAINS; tmp++)
    {
        if (*tmp == '.')
        {
            level_ptr[cur_level] = tmp+1;
            level_length[cur_level-1] = tmp - level_ptr[cur_level-1];
            cur_level++;
        }
    }

    // Too much levels
    if (cur_level >= MAX_SUBDOMAINS)
    {
        fprintf(stderr, "Error: Too much levels for domain %s\n", domain);
        return NULL;
    }

    // Final level
    level_length[cur_level-1] = tmp - level_ptr[cur_level-1];

    tmp = NULL;
    if (cur_level > 2)
    {
        // Standard root domain (zzz.xxx.com) or more
        tmp = level_ptr[1];
    }
    // Simple xxx.com
    else if (cur_level == 2)
        tmp = level_ptr[0];

    if (tmp)
    {
        res = malloc(2+strlen(tmp)+1);
        sprintf(res, "*.%s", tmp);
    }

    return res;
}

static size_t write_callback(char *ptr, size_t size, size_t nmemb, void *userdata)
{
    if ((size*nmemb) > RESPONSE_SIZE)
    {
        fprintf(stderr, "Error curl response is too big (%d bytes, max %d bytes)\n",
                (int)(size*nmemb), RESPONSE_SIZE);
    }
    else
        memcpy(userdata, ptr, size*nmemb);

    return size*nmemb;
}

static int ask_server(struct gpass_parameters* params)
{
    char* wc_domain, *saveptr, *token, *cur_ptr;
    unsigned char* enc_domain;
    unsigned enc_size, matched_key = 0, crypto_v1_index = 1;
    char* request = NULL;
    int ret = -1, res, len;
    CURL *curl;
    char response[RESPONSE_SIZE];
    unsigned char password[256];

    if (params->verbose)
        printf("Username: %s\n", params->username);

    encrypt_domain(params, params->domain, &enc_domain, &enc_size);
    append_to_request(&request, (char*)enc_domain, enc_size);
    free(enc_domain);
    
    wc_domain = wildcard_domain(params->domain);
    if (wc_domain)
    {
        crypto_v1_index++;
        encrypt_domain(params, wc_domain, &enc_domain, &enc_size);
        append_to_request(&request, (char*)enc_domain, enc_size);
        free(enc_domain);
    }

    if (params->crypto_v1_compatible)
    {
        encrypt_domain_v1(params, params->domain, &enc_domain, &enc_size);
        append_to_request(&request, (char*)enc_domain, enc_size);
        free(enc_domain);
        if (wc_domain)
        {
            encrypt_domain_v1(params, wc_domain, &enc_domain, &enc_size);
            append_to_request(&request, (char*)enc_domain, enc_size);
            free(enc_domain);
        }

    }
    
    if (params->verbose)
        printf("Request: %s\n", request);

    curl = curl_easy_init();
    curl_easy_setopt(curl, CURLOPT_URL, params->server);
    curl_easy_setopt(curl, CURLOPT_PORT, params->server_port);
    curl_easy_setopt(curl, CURLOPT_SSL_VERIFYPEER, params->verify_ssl_peer);
    if (params->ca_path)
        curl_easy_setopt(curl, CURLOPT_CAINFO, params->ca_path);
    curl_easy_setopt(curl, CURLOPT_POST, 1);
    curl_easy_setopt(curl, CURLOPT_POSTFIELDS, request);
    if (params->verbose)
        curl_easy_setopt(curl, CURLOPT_VERBOSE, 1);
    curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(curl, CURLOPT_WRITEDATA, (void*)response);

    res = curl_easy_perform(curl);
    curl_easy_cleanup(curl);

    if (res != CURLE_OK)
    {
        fprintf(stderr, "curl_easy_perform() failed: %s\n",
                curl_easy_strerror(res));
        goto end;
    }

    token = strtok_r(response, "\n", &saveptr);

    while (token)
    {
        if (params->verbose)
            printf("Parse %s\n", token);
        cur_ptr = token;
        if (!strcmp(token, "<end>"))
            break;
        else if (!STRNCMP(token, "protocol"))
        {
            cur_ptr += sizeof("protocol"); // includes "="
            if (STRNCMP(cur_ptr, "gpass-"))
            {
                fprintf(stderr, "Error: Unknown server protocol %s\n", token);
                break;
            }
            else
            {
                cur_ptr += sizeof("gpass-")-1;
                if (atoi(cur_ptr) > SERVER_PROTOCOL)
                {
                    fprintf(stderr, "Error: Cannot handle server protocol %s\n", token);
                    break;
                }
            }
        }
        else if (!STRNCMP(token, "pass"))
        {
            cur_ptr += sizeof("pass"); // includes "="

            if ((strlen(cur_ptr)/2) > sizeof(password))
            {
                fprintf(stderr, "Error: retrieved password is too big !\n");
                goto end;
            }

            hex_to_bin(password, (unsigned char*)cur_ptr, strlen(cur_ptr));

            if (matched_key >= crypto_v1_index)
            {
                // Crypto v1
                EVP_DecryptInit(s_cipher_ctx, EVP_aes_256_ecb(), params->derived_master_key, NULL);
                EVP_CipherUpdate(s_cipher_ctx, password, &res, password, strlen(cur_ptr)/2);
                // Remove salt
                password[strlen((char*)password)-3] = 0;
            }
            else
            {
                EVP_DecryptInit(s_cipher_ctx, EVP_aes_256_cbc(), params->derived_master_key, params->global_iv);
                EVP_CipherUpdate(s_cipher_ctx, password, &res, password, strlen(cur_ptr)/2);
                // Remove salt
                len = strlen((char*)password);
                memmove(password, &password[3], len-3);
                password[len-3] = 0;
            }
            display_password((char*)password, DISPLAY_TIME);
            ret = 0;
            goto end;
        }
        else if (!STRNCMP(token, "pbkdf2_level"))
        {
            cur_ptr += sizeof("pbkdf2_level"); // includes "="

            if (atoi(cur_ptr) != params->pbkdf2_level)
            {
                params->pbkdf2_level = atoi(cur_ptr);
                ret = 1;
                break;
            }
        }
        else if (!STRNCMP(token, "matched_key"))
        {
            cur_ptr += sizeof("matched_key"); // includes "="

            matched_key = atoi(cur_ptr);
        }
        else
        {
            fprintf(stderr, "Error: Unknown server response %s\n", token);
            break;
        }
        token = strtok_r(NULL, "\n", &saveptr);
    }

    if (ret)
        printf("Password not found\n");

end:
    free(request);

    return ret;
}

static void init_parameters(struct gpass_parameters* params)
{
    memset (params, 0, sizeof(*params));
    params->pbkdf2_level    = DEFAULT_PBKDF2_LEVEL;
    params->server_port     = DEFAULT_SERVER_PORT;
    params->verify_ssl_peer = 1;
    params->crypto_v1_compatible = 1; // For now, in the next version it must a command line parameter
}

static void release_parameters(struct gpass_parameters* params)
{
    if (params->server)             free(params->server);
    if (params->salt)               free(params->salt);
    if (params->domain)             free(params->domain);
    if (params->username)           free(params->username);
    if (params->orig_master_key)    free(params->orig_master_key);
    if (params->derived_master_key) free(params->derived_master_key);
    if( params->ca_path)            free(params->ca_path);
    if (params->global_iv)          free(params->global_iv);
}

static int check_parameters(struct gpass_parameters* params)
{
    if (!params->server)
    {
        fprintf(stderr, "Error: server not set\n");
        return 1;
    }

    if (!params->domain)
    {
        fprintf(stderr, "Error: gpass domain not set\n");
        return 1;
    }

    if (!params->username)
    {
        fprintf(stderr, "Error: username not set\n");
        return 1;
    }

    return 0;
}

static int gpass_ini_handler(void* user, const char* section,
                       const char* name, const char* value)
{
    struct gpass_parameters* params = (struct gpass_parameters*) user;

    if (!STRNCMP(name, "ca_path"))
    {
        if (params->ca_path) free(params->ca_path);
        params->ca_path = strdup(value);
    }
    else if (!STRNCMP(name, "pbkdf2_level"))
        params->pbkdf2_level = atoi(value);
    else if (!STRNCMP(name, "verify_ssl_peer"))
        params->verify_ssl_peer = atoi(value);
    else if (!STRNCMP(name, "server_port"))
    {
        params->server_port = atoi(value);
        params->port_set = 1;
    }
    else if (!STRNCMP(name, "server"))
    {
        if (params->server) free(params->server);
        params->server = strdup(value);
    }
    else
        fprintf(stderr, "Error: Unknown key '%s' in config file\n", name);

    return 1;
}

static void usage(char* program_name)
{
    fprintf(stderr, "Usage: %s [-f config_file] [-p server_port] [-c CA_certificate_path] [-l PBKDF2_level] [-s gpass_server] [-v] -d domain -u username\n",
            program_name);
    exit(EXIT_FAILURE);
}

int main(int argc, char** argv)
{
    struct gpass_parameters params;
    int opt, ret = 0;
    char* tmp;
    char* config_file, *home;

    if (argc == 1)
        usage(argv[0]);

    init_parameters(&params);

    home = getenv("HOME");
    if (home)
    {
        config_file = malloc(strlen(home)+1+sizeof(DEFAULT_CONFIG_FILE));
        sprintf(config_file, "%s/" DEFAULT_CONFIG_FILE, home);

        ini_parse(config_file, gpass_ini_handler, &params);

        free(config_file);
    }

    while ((opt = getopt(argc, argv, "c:d:f:l:np:s:u:vh")) != -1) {
        switch (opt) {
        case 'c':
            if (params.ca_path) free(params.ca_path);
            params.ca_path = strdup(optarg);
            break;
        case 'd':
            if (params.domain) free(params.domain);
            params.domain = strdup(optarg);
            break;
        case 'f':
            ini_parse(optarg, gpass_ini_handler, &params);
            break;
        case 'l':
            params.pbkdf2_level = atoi(optarg);
            break;
        case 'n':
            params.verify_ssl_peer = 0;
            break;
        case 'p':
            params.server_port = atoi(optarg);
            params.port_set = 1;
            break;
        case 's':
            if (params.server) free(params.server);
            params.server = strdup(optarg);
            break;
        case 'u':
            if (params.username) free(params.username);
            params.username = strdup(optarg);
            break;
        case 'v':
            params.verbose++;
            break;
        case 'h':
        case '?':
        default: /* '?' */
            usage(argv[0]);
        }
    }

    ret = check_parameters(&params);

    if (ret)
        goto end;

    // Manage server, server_port and salt
    if (!STRNCMP(params.server, "http://"))
    {
        if (!params.port_set)
            params.server_port = 80;
        params.salt = strdup(&params.server[7]);
    }
    else if (!STRNCMP(params.server, "https://"))
    {
        if (!params.port_set)
            params.server_port = 443;
        params.salt = strdup(&params.server[8]);
    }

    // Manage domain
    if (!STRNCMP(params.domain, "http://"))
    {
        tmp = strdup(&params.domain[7]);
        free(params.domain);
        params.domain = tmp;
    }
    else if (!STRNCMP(params.domain, "https://"))
    {
        tmp = strdup(&params.domain[8]);
        free(params.domain);
        params.domain = tmp;
    }

    // Remove query part of domain (a.com[/XXXX])
    for (tmp=params.domain; *tmp; tmp++)
    {
        if (*tmp == '/')
        {
            *tmp = 0;
            break;
        }
    }

    s_md_ctx = EVP_MD_CTX_new();
    s_md_256 = EVP_sha256();
    EVP_DigestInit(s_md_ctx, s_md_256);

    s_cipher_ctx = EVP_CIPHER_CTX_new();

    // Let's go
    tmp = getpass("Enter master key: ");

    if (!tmp)
        goto end;

    params.orig_master_key = strdup(tmp);
    derive_master_key(&params);

    // Ctrl+C
    signal(SIGINT, signal_handler);
    // Ctrl+Z
    signal(SIGTSTP, signal_handler);

    ret = ask_server(&params);

    // try again with new parameters
    if (ret > 0)
    {
        derive_master_key(&params);
        ask_server(&params);
    }
    
end:
    release_parameters(&params);

    if (s_md_ctx) EVP_MD_CTX_free(s_md_ctx);
    if (s_cipher_ctx) EVP_CIPHER_CTX_free(s_cipher_ctx);

    return ret;
}