/*
  Copyright (C) 2013-2016 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 <curl/curl.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_PKDBF2_LEVEL    1000
#define MASTER_KEY_LENGTH    (256/8)
#define BLOCK_SIZE           (128/8)
#define DEFAULT_SERVER_PORT      443
#define SERVER_PROTOCOL            3
#define RESPONSE_SIZE           2048
#define MAX_SUBDOMAINS            10

struct gpass_parameters {
    unsigned pkdbf2_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;
} ;

static void derive_master_key(struct gpass_parameters* params)
{
    if (!params->derived_master_key)
        params->derived_master_key = malloc(MASTER_KEY_LENGTH);
    
    PKCS5_PBKDF2_HMAC(params->orig_master_key, strlen(params->orig_master_key),
                      (unsigned char*)params->salt, strlen(params->salt),
                      params->pkdbf2_level, EVP_sha256(),
                      MASTER_KEY_LENGTH, params->derived_master_key);
}

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(struct gpass_parameters* params, char* domain,
                           unsigned char** res, unsigned* out_size)
{
    EVP_CIPHER_CTX* evp_ctx;
    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_ctx = EVP_CIPHER_CTX_new();
    EVP_EncryptInit(evp_ctx, EVP_aes_256_ecb(), params->derived_master_key, NULL);
    EVP_CipherUpdate(evp_ctx, tmp, (int*)out_size, buffer, size);
    EVP_CIPHER_CTX_free(evp_ctx);

    bin_to_hex(tmp, *res, size);

    *out_size *= 2;

    free(buffer);
    free(tmp);
}

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

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

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

        sprintf(&((*request)[strlen(*request)]), "&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] = tmp - level_ptr[cur_level-1];

    tmp = NULL;
    if (cur_level >= 3)
    {
        // Seems to be a two level root domain (ie zzz.xxx.co.uk ...)
        if (level_length[cur_level-2] <= 3)
        {
            if (cur_level > 3)
                tmp = level_ptr[cur_level-3];
        }
        else
            // Standard root domain (zzz.xxx.com)
            tmp = level_ptr[cur_level-2];
    }
    // 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;
    char* request = NULL;
    int ret = -1, res;
    CURL *curl;
    EVP_CIPHER_CTX* evp_ctx;
    char response[RESPONSE_SIZE];
    unsigned char password[256];

    encrypt_domain(params, params->domain, &enc_domain, &enc_size);
    append_to_request(&request, (char*)enc_domain);
    free(enc_domain);

    wc_domain = wildcard_domain(params->domain);
    if (wc_domain)
    {
        encrypt_domain(params, wc_domain, &enc_domain, &enc_size);
        append_to_request(&request, (char*)enc_domain);
        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));

            evp_ctx = EVP_CIPHER_CTX_new();
            EVP_DecryptInit(evp_ctx, EVP_aes_256_ecb(), params->derived_master_key, NULL);
            EVP_CipherUpdate(evp_ctx, password, &res, password, strlen(cur_ptr)/2);
            EVP_CIPHER_CTX_free(evp_ctx);

            // Remove salt
            password[strlen((char*)password)-3] = 0;
            printf("Password found: %s\n", password);
            ret = 0;
            goto end;
        }
        else if (!STRNCMP(token, "pkdbf2_level"))
        {
            cur_ptr += sizeof("pkdbf2_level"); // includes "="

            if (atoi(cur_ptr) != params->pkdbf2_level)
            {
                params->pkdbf2_level = atoi(cur_ptr);
                ret = 1;
                break;
            }
        }
        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->pkdbf2_level    = DEFAULT_PKDBF2_LEVEL;
    params->server_port     = DEFAULT_SERVER_PORT;
    params->verify_ssl_peer = 1;
}

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);
}

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, "pkdbf2_level"))
        params->pkdbf2_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 PKDBF2_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.pkdbf2_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;
        }
    }

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

    if (!tmp)
        goto end;

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

    ret = ask_server(&params);

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

end:
    release_parameters(&params);

    return ret;
}