#include "member-tpm.h"

typedef enum memberstate {
    ON,
    ISSUERPUB,
    RCVPUBLIC,
    JOIN,
    APPEND,
    JOINPROCEED,
    JOINED,
} memberstate_e;

typedef struct member {
    struct ecdaa_member_public_key_FP256BN mpk;
    memberstate_e state;
    uint8_t nonce[NONCE_SIZE];
    struct ecdaa_credential_FP256BN cred;
    struct ecdaa_issuer_public_key_FP256BN ipk;
    uint8_t bsn[MAX_BSNSIZE];
    size_t bsn_len;
    struct ecdaa_tpm_context ctx;
    TPM2_HANDLE pk_handle;
    unsigned char pk_in[ECP_FP256BN_LENGTH];
    unsigned char tcti_buffer[256];
} member_t;

member_t member;
uint8_t msg[MAX_MSGSIZE];
size_t msg_len;
uint8_t chksum[MAX_CHKSUMSIZE];
size_t chksum_len;

int init_tpm();
int free_tpm();

static int read_public_key_from_files(uint8_t *public_key, TPM2_HANDLE *key_handle, const char *tpm_key_file, const char *tpm_handle_file);

int member_join(char *buffer);

int member_verifymsg(char *buffer);

int member_publish(char *buffer);

int member_joinappend(char *buffer);

int member_joinfinish(char *buffer);

int main(int argc, char *argv[]) {
    char buffer[MAX_BUFSIZE];
    char *remote_ip;
    int ret = 0;
    TPM2_HANDLE sk_handle = 0;

    //strncpy(member.bsn, "mybasename", 10);
    //member.bsn_len = strlen(member.bsn);
    switch(argc) {
        case 3:
            if(0 == strncasecmp("--join", argv[1], 6) || 0 == strncasecmp("-j", argv[1], 2)) {
		if (0 != read_public_key_from_files(member.pk_in, &sk_handle, MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE)) {
                    printf("Could not load TPM key with '%s' and '%s', trying to create a new key...\n", MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE);
		    if (0 != create_key(MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE) || 0 != read_public_key_from_files(member.pk_in, &sk_handle, MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE)) {
                        printf("Error: Creating or Loading TPM key with '%s' and '%s' failed.\n", MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE);
                        return 1;
                    }
		}
                if (0 != init_tpm()) {
                    printf("Error: Failed to initialize TPM.\n");
                    return 1;
                }
                printf("Initialized TPM with pubkey and handle.\n");
                member.state = ON;
                remote_ip = argv[2];
                ret = client_connect(&member_join, remote_ip, ISSUERPORT);
                if (0 >= ret || JOINED != member.state) {
                    printf("Join process failed!\n");
                    return 1;
                }
                printf("Join process was successful.\n");
            } else if (0 == strncasecmp("--send", argv[1], 6) || 0 == strncasecmp("-s", argv[1], 2)) {
		if (0 != read_public_key_from_files(member.pk_in, &sk_handle, MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE)) {
                    printf("Error: reading in public key files '%s' and '%s' failed\n", MEMBER_TPM_KEY_FILE, MEMBER_TPM_HANDLE_FILE);
		    return 1;
		}
                if (0 != init_tpm()) {
                    printf("Error: Failed to initialize TPM.\n");
                    return 1;
                }
                printf("Initialized TPM with pubkey and handle.\n");
                msg_len = ecdaa_read_from_file(msg, MAX_MSGSIZE, MESSAGE_FILE);
		if (msg_len < 0) {
                    printf("Could not open message file %s.\n", MESSAGE_FILE);
                    return 1;
		}
                chksum_len = ecdaa_read_from_file(chksum, MAX_CHKSUMSIZE, CHECKSUM_FILE);
		if (chksum_len < 0) {
                    printf("Could not open checksum file %s.\n", CHECKSUM_FILE);
                    return 1;
		}
                printf("Loaded message and checksum.\n");
                if (0 > ecdaa_read_from_file(member.nonce, NONCE_SIZE, MEMBER_NONCE_FILE) ||
                        0 != ecdaa_member_public_key_FP256BN_deserialize_file(&member.mpk, MEMBER_PUBLIC_KEY_FILE, member.nonce, NONCE_SIZE) ||
                        0 != ecdaa_credential_FP256BN_deserialize_file(&member.cred, MEMBER_CREDENTIAL_FILE)) {
                    printf("Could not import key files. Importing from either %s, %s or %s was not successful.\n",
                        MEMBER_NONCE_FILE, MEMBER_PUBLIC_KEY_FILE, MEMBER_CREDENTIAL_FILE);
                    return 1;
                }
                member.state = JOINED;
                remote_ip = argv[2];
                ret = client_connect(&member_verifymsg, remote_ip, VERIFIERPORT);
                if (2 != ret) {
                    printf("Error: Message transmission to verifier failed.\n");
                }
            } else {
                printf("Error: Arguments invalid.\n");
            }
            break;
        default:
            printf("Usage: \n Join an issuer's group: %s --join <issuer's IPv4>\n", argv[0]);
            printf("Send a signed message to the verifier: %s --send <verifier's IPv4>\n", argv[0]);
            printf("Before sending a DAA-signed message, the member must join a DAA group\n");
            printf("%s must not exceed %d Bytes, %s must be smaller than %d Bytes\n", MESSAGE_FILE, MAX_MSGSIZE, CHECKSUM_FILE, MAX_CHKSUMSIZE);
            break;
    }
    return 0;
}

int init_tpm()
{
    const char *device_conf = "/dev/tpm0";

    int ret = 0;

    memset(member.tcti_buffer, 0, sizeof(member.tcti_buffer));

    TPM2_HANDLE key_handle = 0;

    TSS2_TCTI_CONTEXT *tcti_ctx = (TSS2_TCTI_CONTEXT*)member.tcti_buffer;

    size_t size;
    ret = Tss2_Tcti_Device_Init(NULL, &size, device_conf);
    if (TSS2_RC_SUCCESS != ret) {
        printf("Error: Failed to get allocation size for tcti context.\n");
        return -1;
    }
    if (size > sizeof(member.tcti_buffer)) {
        printf("Error: device TCTI context size larger than pre-allocated buffer.\n");
        return -1;
    }
    ret = Tss2_Tcti_Device_Init(tcti_ctx, &size, device_conf);
    if (TSS2_RC_SUCCESS != ret) {
        printf("Error: Unable to initialize device TCTI context.\n");
        return -1;
    }

    //initialize ecdaa tpm context
    if(0 != ecdaa_tpm_context_init(&member.ctx, member.pk_handle, NULL, 0, tcti_ctx)) {
        printf("Error: ECDAA context failed.\n");
        return -1;
    }
#ifdef DEBUG
    printf("ECDAA context initialized.\n");
#endif
    return 0;
}

int free_tpm() {
    ecdaa_tpm_context_free(&member.ctx);
    return 0;
}

int member_join(char *buffer) {
    int ret = 0;

    switch (member.state) {
        case ON:
            bzero(buffer, MAX_BUFSIZE);
            strncpy(buffer, "PUBLISH\n", 8);
            member.state = ISSUERPUB;
            break;
        case ISSUERPUB:
            if (0 == strncasecmp("PUBLISH", buffer, 7)) {
#ifdef DEBUG
                printf("ISSUER > MEMBER: %s\n", buffer);
#endif
                uint8_t binbuf[MAX_BUFSIZE];
                char *current = &buffer[8];
                ecdaa_decode(current, binbuf, ECDAA_ISSUER_PUBLIC_KEY_FP256BN_LENGTH);
                ret = ecdaa_issuer_public_key_FP256BN_deserialize(&member.ipk, binbuf);
                if (-1 == ret) {
                    printf("member_getpublic: issuer public key is malformed!\n");
                    ret = -1;
                } else if (-2 == ret) {
                    printf("member_getpublic: signature of issuer public key is invalid.\n");
                    ret = -1;
                } else {
                    bzero(buffer, MAX_BUFSIZE);
                    strncpy(buffer, "JOIN\n", 5);
                    member.state = APPEND;
                    ret = 0;
                }
            } else {
                printf("member_getpublic: did not get public key from issuer.\n");
                member.state = ON;
                ret = -1;
            }
            break;
        case APPEND:
            if (0 == strncasecmp("JOINSTART", buffer, 9)) {
#ifdef DEBUG
                printf("ISSUER > MEMBER: %s\n", buffer);
#endif
                member_joinappend(buffer);
                member.state = JOINPROCEED;
            } else {
                printf("member_join: did not get nonce from issuer.\n");
                member.state = RCVPUBLIC;
                ret = -1;
            }
            break;
        case JOINPROCEED:
            if (0 == strncasecmp("JOINPROCEED", buffer, 11)) {
#ifdef DEBUG
                printf("ISSUER > MEMBER: %s\n", buffer);
#endif
                ret = member_joinfinish(buffer);
                member.state = JOINED;
            } else {
                printf("member_getpublic: did not get credentials from issuer.\n");
                member.state = RCVPUBLIC;
                ret = -1;
            }
            break;
        default:
            ret = -1;
    }
    if (0 == ret) {
#ifdef DEBUG
        printf("ISSUER < MEMBER: %s", buffer);
#endif
    }
    return ret;
}

//"VERIFYMSG <msg><checksum>0<signature>" or
//"VERIFYMSG <msg><checksum>1<signature with bsn>"
int member_verifymsg(char *buffer) {
    char *current = buffer;
    uint8_t binbuf[MAX_BUFSIZE];
    uint8_t has_nym = member.bsn_len > 0 ? 1 : 0;
    struct ecdaa_signature_FP256BN sig;
    size_t sig_len = has_nym ? ecdaa_signature_FP256BN_with_nym_length() : ecdaa_signature_FP256BN_length();
    int bytes = 0;

    if (0 == strncasecmp("OK", buffer, 2)) {
        return 1;
    } else if (0 == strncasecmp("ERR", buffer, 3)) {
        printf("member_verifymsg: Verifier refused signature.\n");
        return 1;
    }

    bzero(buffer, MAX_BUFSIZE);
    strncpy(current, "VERIFYMSG ", 10);
    current = &current[10];

    bytes = ecdaa_encode(msg, current, MAX_MSGSIZE);
    current = &current[bytes];
    bytes = ecdaa_encode(chksum, current, MAX_CHKSUMSIZE);
    current = &current[bytes];
    if(has_nym) {
        if (0 != ecdaa_signature_TPM_FP256BN_sign(&sig, chksum, chksum_len, member.bsn, member.bsn_len, &member.cred, ecdaa_rand, &member.ctx)) {
            printf("member_verifymsg: Signing message failed.\n");
        }
        current[0] = '1';
        current = &current[1];
        strncpy(current, (char *)member.bsn, MAX_BSNSIZE);
        current = &current[MAX_BSNSIZE];
    } else {
        if (0 != ecdaa_signature_TPM_FP256BN_sign(&sig, chksum, chksum_len, NULL, 0, &member.cred, ecdaa_rand, &member.ctx)) {
            printf("member_verifymsg: Signing message failed.\n");
        }
        current[0] = '0';
        current = &current[1];
    }

    bzero(binbuf, MAX_BUFSIZE);
    ecdaa_signature_FP256BN_serialize(binbuf, &sig, has_nym);
    bytes = ecdaa_encode(binbuf, current, sig_len);
#ifdef DEBUG
    printf("member_verifymsg: has_nym: %u, sig_len: %lu\n",has_nym, sig_len);
    printf("member_verifymsg: msg: %s, len: %lu\n", msg, msg_len);
    printf("member_verifymsg: chksum: %s, len: %lu\n", chksum, chksum_len);
    printf("member_verifymsg: bsn: %s, len: %lu\n", (char *)member.bsn, strlen((char *)member.bsn));
    printf("member_verifymsg: sig: %s, len: %lu\n", current, sig_len);
#endif
    current[bytes] = '\n';
    return 2; //send to verifier before closing
}

//"JOINSTART <issuer.nonce>" > "APPEND <member.mpk>"
int member_joinappend(char *buffer) {
    char *current = &buffer[10];
    uint8_t binbuf[MAX_BUFSIZE];
    int ret = 0;
    int bytes = ecdaa_decode(current, member.nonce, NONCE_SIZE);
    ecdaa_write_buffer_to_file(MEMBER_NONCE_FILE, member.nonce, NONCE_SIZE);
    bzero(buffer, MAX_BUFSIZE);
    if (0 != (ret = ecdaa_member_key_pair_TPM_FP256BN_generate(&member.mpk, member.pk_in, member.nonce, NONCE_SIZE, &member.ctx))) {
        fprintf(stderr, "Error generating member key-pair: ret = %d\n", ret);
        strncpy(buffer, "ERR\n", 4);
        return -1;
    }
    strncpy(buffer, "APPEND ", 7);

    current = &buffer[7];
    bzero(binbuf, MAX_BUFSIZE);
    ecdaa_member_public_key_FP256BN_serialize(binbuf, &member.mpk);
    bytes = ecdaa_encode(binbuf, current, ECDAA_MEMBER_PUBLIC_KEY_FP256BN_LENGTH);
    current[bytes] = '\n';
    return 0;
}

//"JOINPROCEED <member.cred><member.cred_sig>" > ""
int member_joinfinish(char *buffer) {
    char *current = &buffer[12];
    uint8_t *bincur;
    uint8_t binbuf[MAX_BUFSIZE];
    int bytes = 0;
    int ret = 0;
    bzero(binbuf, MAX_BUFSIZE);
    bytes = ecdaa_decode(current, binbuf, ECDAA_CREDENTIAL_FP256BN_LENGTH);

    current = &current[bytes];
    bincur = &binbuf[ECDAA_CREDENTIAL_FP256BN_LENGTH];
    ecdaa_decode(current, bincur, ECDAA_CREDENTIAL_FP256BN_SIGNATURE_LENGTH);
    ret = ecdaa_credential_FP256BN_deserialize_with_signature(&member.cred, &member.mpk, &member.ipk.gpk, binbuf, bincur);
    if(-1 == ret) {
        printf("member_joinfinish: credential is malformed!\n");
        ret = -1;
    } else if(-2 == ret) {
        printf("member_joinfinish: siganture of credential is invalid.\n");
        ret = -1;
    }
    printf("member_joinfinish: writing public key and credential to disk.\n");
    if(0 != ecdaa_member_public_key_FP256BN_serialize_file(MEMBER_PUBLIC_KEY_FILE, &member.mpk) ||
            0 != ecdaa_credential_FP256BN_serialize_file(MEMBER_CREDENTIAL_FILE, &member.cred)) {
        printf("issuer_setup: Error saving key-pair or credential to disk.\n");
        ret = -1;
    } else {
        ret = 1;
    }

    return ret;
}

int read_public_key_from_files(uint8_t *public_key, TPM2_HANDLE *key_handle, const char *pub_key_filename, const char *handle_filename)
{
    int ret = 0;

    FILE *pub_key_file_ptr = fopen(pub_key_filename, "r");
    if (NULL == pub_key_file_ptr) {
    	printf("read_public_key: error opening public key.\n");
        return -1;
    }
    do {
        for (unsigned i=0; i < ECP_FP256BN_LENGTH; i++) {
            unsigned byt;
            if (ret = fscanf(pub_key_file_ptr, "%02X", &byt) != 1) {
                ret = -1;
                break;
            }
            public_key[i] = (uint8_t)byt;
        }
    } while(0);
    (void)fclose(pub_key_file_ptr);
    if (0 != ret) {
    	printf("read_public_key: error reading public key file.\n");
        return -1;
    }

    FILE *handle_file_ptr = fopen(handle_filename, "r");
    if (NULL == handle_file_ptr) {
    	printf("read_public_key: error opening handle.\n");
        return -1;
    }

    do {
        for (int i=(sizeof(TPM2_HANDLE)-1); i >= 0; i--) {
            unsigned byt;
            if (fscanf(handle_file_ptr, "%02X", &byt) != 1) {
                ret = -1;
                break;
            }
            *key_handle += byt<<(i*8);
        }
        if (0 != ret)
            break;
    } while(0);
    (void)fclose(handle_file_ptr);

    if (0 != ret) {
    	printf("read_public_key: error closing public key.\n");
        return -1;
    }
    return ret;
}