// Copyright (c) 2019  Diamond Key Security, NFP
// Copyright (c) 2019  Diamond Key Security, NFP
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
// - Redistributions of source code must retain the above copyright notice,
//   this list of conditions and the following disclaimer.
//
// - Redistributions in binary form must reproduce the above copyright
//   notice, this list of conditions and the following disclaimer in the
//   documentation and/or other materials provided with the distribution.
//
// - Neither the name of the Diamond Key Security nor the names of its contributors may
//   be used to endorse or promote products derived from this software
//   without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS “AS
// IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
// TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
// PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
// TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <string.h>
#include <termios.h>

#include <djson.h>

#include "cryptech_device.h"
#include "cryptech_device_cty.h"

// Internal Function Declarations ------------------------------------------
void ResetKeyboardInput(struct termios *oldAttributes);
void SetRawKeyboardInput(struct termios *oldAttributes);
char *GetPassword(char *buffer, int buffer_len);
int GetOption(const char *banner, const char *answers, const char *question);
void GetMasterKey(char *buffer, int buffer_len);
void GetLine(char *buffer, int buffer_len);
int GetLineCheck(char *buffer, int buffer_len, const char *question);
int isMasterKeyValid(char *buffer);
int SetMasterKey(char *masterkey, char *pin);

void SaveSetupJson(FILE **fp, uint32_t handle);
void SaveExportJson(FILE **ofp, char *setup_json, uint32_t handle);
void ImportKeys(char *import_json, uint32_t handle);

// Internal Enumerations --------------------------------------------------
typedef enum
{
    cmd_op_setup = 0,
    cmd_op_export = 1,
    cmd_op_import = 2
} backup_operations_t;

// Function Definintions --------------------------------------------------
int main(int argc, char *argv[])
{
    char buffer[256];
    char pin[64];
    char masterkey[128];
    char inputfile[2048];
    char outputfile[2048];
    inputfile[0] = 0;
    outputfile[0] = 0;
    FILE *ofp = NULL;
    char *input_json = NULL;

    printf("dks_cryptech_backup\r\nCopyright 2019 Diamond Key Security, NFP.\r\n\
Port of cryptech_backup from CrypTech with support for copying key attributes.\r\n\r\n\
Uses code from the CrypTech project\r\n\
Copyright (c) 2016-2018, NORDUnet A/S\r\n\
This works by having the destination HSM (the one importing keys)\r\n\
create an RSA keypair (the \"KEKEK\"), the public key of which can then\r\n\
be imported into the source HSM (the one exporting keys) and used to\r\n\
encrypt AES key encryption keys (KEKs) which in turn can be used to\r\n\
wrap the private keys being transfered.  Transfers are encoded in\r\n\
JSON; the underlying ASN.1 formats are SubjectPublicKeyInfo (KEKEK\r\n\
public key) and PKCS #8 EncryptedPrivateKeyInfo (everything else).\r\n\r\n\
NOTE WELL: while this process makes it POSSIBLE to back up keys\r\n\
securely, it is not sufficient by itself: the operator MUST make\r\n\
sure only to export keys using a KEKEK known to have been generated by\r\n\
the target HSM.\r\n");

    printf("\r\nversion 19.07\r\n\r\n");

    printf("\r\n\r\n----------------------------------------------------------\r\n");
    printf("\r\nPlease enter the 'wheel' password for the CrypTech device.\r\n");
    GetPassword(pin, 64);

    int setmasterkey = GetOption(
"Would you like to set the master key on the Cryptech device?\r\n\
  Y) Yes\r\n\
  N) No\r\n\
  Q) Quit\r\n", "YyNnQq", "Please select an option (Y, N, Q): ");
    if (setmasterkey == 2) return 0;
    if (setmasterkey == 0) GetMasterKey(masterkey, 80);

    int mode = GetOption("Please select a backup operation.\r\n\
  S) Setup - Create a KEKEK on the CrypTech device and save a 'setup.json' file.\r\n\
  E) Export - Load a KEKEK from an external device from a 'setup.json' file and save a 'export.json' file.\r\n\
  I) Import - Import data from an 'export.json' file.\r\n\
  Q) Quit\r\n", "SsEeIiQq", "Please choose a backup operation (S, E, I, Q): ");
    if (mode == 3) return 0;

    if (mode == cmd_op_export || mode == cmd_op_import)
    {
        if(GetLineCheck(inputfile, sizeof(inputfile)/sizeof(char),
                        "\r\nPlease enter the file path of the input file:\r\n> ") == 0) return 0;
    }
    if (mode == cmd_op_setup || mode == cmd_op_export)
    {
        if(GetLineCheck(outputfile, sizeof(outputfile)/sizeof(char),
                        "\r\nPlease enter the file path of the output file:\r\n> ") == 0) return 0;
    }

    const char *mode_strings[] = { "Setup", "Export", "Import"};

    printf("\r\n\r\n----------------------------------------------------------\r\n");
    printf("Please confirm options:\r\n");
    printf("  Mode:%s\r\n", mode_strings[mode]);
    printf("  Master key: ");
    if(setmasterkey == 1) printf("The master key has already been set on the device.\r\n");
    else if (strlen(masterkey) == 0) printf("A random key will be generated.\r\n");
    else printf("%s\r\n", masterkey);
    if(inputfile[0] != 0) printf("  Input file: %s\r\n", inputfile);
    if(outputfile[0] != 0) printf("  Output file: %s\r\n", outputfile);

    if (GetOption("\r\nContinue with these options?\r\n\
  Y) Yes\r\n\
  N) No\r\n", "NnYy", "Please select an option (Y, N): ") == 0) return 0;

    // try to open files
    if(inputfile[0] != 0)
    {
        input_json = djson_loadfile(inputfile);
        if(input_json == NULL)
        {
            printf("\r\nUnable to open input file, '%s'.\r\n", inputfile);
            goto done;
        }
    }

    if(outputfile[0] != 0)
    {
        ofp = fopen(outputfile, "wt");
        if(ofp == NULL)
        {
            printf("\r\nUnable to open output file, '%s'.\r\n", outputfile);
            goto done;
        }
    }

    if(setmasterkey == 0)
    {
        printf("\r\nSetting the master key on the CrypTech device.\r\n");
        if (SetMasterKey(masterkey, pin) != 0)
        {
            printf("\r\nUnable to set the master key.\r\n");
            goto done;
        }
        printf("\r\nMaster key set.\r\n\r\n");
    }

    uint32_t handle = get_random_handle();

    int rval = init_cryptech_device(pin, handle);

    if (rval != 0)
    {
        printf("Unable to log into CrypTech device.\r\n");
        printf("Did you run 'eval $(cryptech_probe)' before running this command?\r\n");
    }
    else
    {
        if (mode == cmd_op_setup)
        {
            SaveSetupJson(&ofp, handle);
        }
        else if (mode == cmd_op_export)
        {
            SaveExportJson(&ofp, input_json, handle);
        }
        else if (mode == cmd_op_import)
        {
            ImportKeys(input_json, handle);
        }

        close_cryptech_device(handle);  
    }

done:
    free(input_json);
    if (ofp != NULL) fclose(ofp);
    return 0;
}

void SaveSetupJson(FILE **fp, uint32_t handle)
{
    printf("\r\nGenerating setup json with KEKEK.\r\n");

    char *setup_json;
    int rval = setup_backup_destination(handle, -1, &setup_json);

    if (rval == 0)
    {
        char *s = setup_json;
        while (*s != 0)
        {
            fputc(*s++, *fp);
        }
        fclose(*fp);
        *fp = NULL;

        printf("KEKEK written to output json.\r\n");
    }
    else
    {
        printf("Unable to get KEKEK.\r\nPlease check that the master key is correct.\r\n");
    }

    // something went wrong
    return;
}

void SaveExportJson(FILE **export_json, char *setup_json, uint32_t handle)
{
    int rval = cryptech_export_keys(handle, setup_json, export_json);
    if (rval == 0)
    {
        // keep export data in a temporary file that we can access when the
        // HSM ask for it
        fclose(*export_json);
        *export_json = NULL;
    }
    else
    {
        printf("\r\nFailure:%i, exporting data.\r\n", rval);
    }
    
}

void ImportKeys(char *import_json, uint32_t handle)
{
    int rval = import_keys(handle, import_json);

    if (rval == 0)
    {
        printf("Import complete\r\n");
    }
    else
    {
        printf("Unable to import data into CrypTech device.\r\n");
    }
}

int SetMasterKey(char *masterkey, char *pin)
{
    int rval = open_cryptech_device_cty();

    if (rval == 0)
    {
        if (strlen(masterkey) == 0)
        {
            // random master key
            rval = cty_setmasterkey(pin, NULL);
        }
        else
        {       
            // user has a master key to use
            rval = cty_setmasterkey(pin, masterkey);
        }

        // make sure we've logged out
        cty_logout();

        rval = close_cryptech_device_cty();
    }

    return rval;
}

void GetMasterKey(char *buffer, int buffer_len)
{
    while (1)
    {
        printf("\r\nPlease enter the master key or leave it blank to set it to a random value\r\n\
  -------- -------- -------- -------- -------- -------- -------- --------\r\n\
> ");
        GetLine(buffer, buffer_len);

        if(strlen(buffer) == 0 || isMasterKeyValid(buffer)) return;
        else
        {
            printf("\r\nInvalid master key. Please try again.\r\n\r\n");
        }
    }
}

int isMasterKeyValid(char *buffer)
{
    int count = 0;
    char *c = buffer;

    while (*c != 0)
    {
        if (*c != ' ')
        {
            ++count;
            if ((*c < 'a' || *c > 'f') &&
                (*c < 'A' || *c > 'F') &&
                (*c < '0' || *c > '9'))
            {
                return 0;
            }
        }
        ++c;
    }

    return 1;
}

int GetLineCheck(char *buffer, int buffer_len, const char *question)
{
    int result;

    while(1)
    {
        printf("%s", question);

        GetLine(buffer, buffer_len);

        if (strlen(buffer) < 1) continue;

        printf("Is '%s' correct?", buffer);

        result = GetOption("  Y) Yes\r\n\
  N) No\r\n\
  Q) Quit\r\n", "NnYyQq", "Please select an option (Y, N, Q): ");

        if (result == 1) return 1;
        else if (result == 2) return 0;
    }
}

void GetLine(char *buffer, int buffer_len)
{
    int i = 0;
    char c;
    while ((c = getchar()) != '\n' && (c != '\r') && (i < buffer_len-1))
    {
        if ((c == 127 || c == 8) && i > 0) --i;
        else buffer[i++] = c;
    }

    // clear stdin incase the user went over the size of the buffer
    if (c != '\n' && c != '\r') while ((c = getchar()) != '\n' && c != '\r');

    buffer[i] = 0;
}

int GetOption(const char *banner, const char *answers, const char *question)
{
    // use big buffers to cover user errors
    char buffer[64];

    printf("\r\n%s", banner);
    while (1)
    {
        printf("\r\n%s", question);

        GetLine(buffer, sizeof(buffer)/sizeof(char));
        if(strlen(buffer) != 1) continue;

        char user_input = buffer[0];

        int i = 0;
        while (answers[i] != 0)
        {
            if(answers[i] == user_input) return i/2;
            ++i;
        }
    }
}

char *GetPassword(char *buffer, int buffer_len)
{
    struct termios originalSTDINAttr;
    printf("Password: ");
    SetRawKeyboardInput(&originalSTDINAttr);

    GetLine(buffer, buffer_len);

    ResetKeyboardInput(&originalSTDINAttr);

    return buffer;
}

// set keyboard input to be non-buffered and do not echo
void SetRawKeyboardInput(struct termios *oldAttributes)
{
    struct termios attr;

    // get a copy of the current settings so we can reset later
    tcgetattr(STDIN_FILENO, oldAttributes);

    // get another copy for setting the new state
    tcgetattr(STDIN_FILENO, &attr);

    // make raw
    cfmakeraw(&attr);
	
    // set
    tcsetattr(STDIN_FILENO, TCSANOW, &attr);
}

// set stdin to saved attributes
void ResetKeyboardInput(struct termios *oldAttributes)
{
    tcsetattr(STDIN_FILENO, TCSANOW, oldAttributes);
}
