/*
 * =====================================================================================
 *
 *       Filename:  tz_pebble_interface.h
 *
 *    Description:  Commom interface for pebble
 *
 *        Version:  1.0
 *        Created:  06/02/2020
 *       Revision:  none
 *       Compiler:  gcc
 *
 *        Company:  Samsung Electronics
 *        Copyright (c) 2020 by Samsung Electronics, All rights reserved.
 *
 * =====================================================================================
 */

#ifndef _TZ_PEBBLE_INTERFACE_H_
#define _TZ_PEBBLE_INTERFACE_H_

/**
 * Includes
 */
#include "stdint.h"
#include "tz_msg.h"
#include <openssl/sha.h>

/**
 * Length
 */
#define CERT_MAX_LEN                2500
#define CERTS_MAX_NUM               3
#define JWS_LEN                     8192
#define NWS_INFO_LEN                128
#define PEBBLE_DRK_MAX_BUF_LEN      4096
#define SERVICE_NAME_LEN            32
#define SSS_KEY_MAX_LEN             32
#define SSS_SHARE_MAX_LEN           SSS_KEY_MAX_LEN + 1
#define ATN_NONCE_MIN_LEN           8
#define ATN_NONCE_MAX_LEN           32
#define KEY_ALIAS_LEN               256

#define SK_MASTER_KEY_LEN 32
#define SK_SUBKEY_LEN 32
#define SK_KEY_MAX_LEN (SK_MASTER_KEY_LEN > SK_SUBKEY_LEN ? SK_MASTER_KEY_LEN : SK_SUBKEY_LEN)
#define SK_WRAP_KEY_LEN 256
#define WRAPPING_AES_KEY_LEN 32
#define WRAPPING_AES_IV_LEN 12
#define WRAPPING_AES_AAD_LEN 54
#define WRAPPING_AES_TAG_LEN 16
#define WRAPPING_RSA_N_BYTES_LEN 256 // In RSA-2048, n's length is 256 bytes
//#define WRAPPING_RSA_N_STR_LEN 617 // RSA-2048 has 617 decimal digits (2,048 bits)
#define RSA_F_RSA_EAY_PUBLIC_ENCRYPT 104

#define MAX_BUDDIES                 5
#define CHECKSUM_LEN                SHA256_DIGEST_LENGTH
#define JWE_JWS_MAX_SIZE 			4096
#define KEY_METADATA_MAX_LEN  		1024
#define SHARE_PAYLOAD_LEN 			SSS_SHARE_MAX_LEN + CHECKSUM_LEN * MAX_BUDDIES + 20 + KEY_METADATA_MAX_LEN

/**
 * PEBBLE Commands
 */
#define CMD_LOAD_CERT               0x00000001
#define CMD_GEN_ATN_NONCE           0x00000002
#define CMD_GEN_MASTER_KEY          0x00000003
#define CMD_DERIVE_SUBKEY           0x00000004
#define CMD_GEN_BACKUP_REQ          0x00000005
#define CMD_GEN_BACKUP_ACK          0x00000006
#define CMD_GEN_RECOVERY_REQ        0x00000007
#define CMD_GEN_RECOVERY_ACK        0x00000008
#define CMD_HANDLE_RECOVERY_ACK     0x00000009
#define CMD_RECONSTRUCT_KEY         0x0000000A
//#define CMD_SSS_SPLIT               0x00000101
//#define CMD_SSS_JOIN                0x00000102
//#define CMD_GET_DEVICE_ID           0x00000000

#define BUF_T(MAX_BUF_LEN) \
	struct { \
	   uint32_t len; \
	   uint8_t buf[MAX_BUF_LEN]; \
	} __attribute__ ((packed))

#define DECLARE_BUF_T(MAX_BUF_LEN) \
	typedef struct { \
	   uint32_t len; \
	   uint8_t buf[MAX_BUF_LEN]; \
	} __attribute__ ((packed))

typedef struct {
	uint32_t len;
	uint8_t blob[CERT_MAX_LEN];
} __attribute__ ((packed)) cert_t;

/**
 *  * DRK certificates chain
 *   */
typedef struct {
	cert_t cert[CERTS_MAX_NUM];
	uint32_t num_certs;
} __attribute__ ((packed)) cert_chain_t;

/**
 * Possible return codes
 */
typedef enum pebble_return_code {
	PEBBLE_STATUS_SUCCESS = 0x0,
	PEBBLE_STATUS_FAIL = 0x1,
	PEBBLE_INVALID_JWS = 0x2,
	PEBBLE_JWS_INVALID_FORMAT = 0x3,
	PEBBLE_JWS_NULL = 0x4,
	PEBBLE_JWS_INVALID_LENGTH = 0x5,
	PEBBLE_INVALID_SIGNATURE = 0x6,
	PEBBLE_INVALID_CERTCHAIN = 0x7,
	PEBBLE_INVALID_CA = 0x8,
	PEBBLE_COMPUTE_HASH_FAIL = 0xC,
	PEBBLE_GEN_DEVICE_ID_FAIL = 0x10,
	PEBBLE_KEY_ERROR = 0x11,
	PEBBLE_SIGNATURE_ERROR = 0x12,
	PEBBLE_CERT_DRK_FAIL = 0x13,
	PEBBLE_ALLOC_ERROR = 0x14,
	PEBBLE_GEN_RESPONSE_FAIL = 0x15,
	PEBBLE_BUF_SIZE_ERROR = 0x16,
	PEBBLE_CONVERT_DER_CERT_FAIL = 0x17,
	PEBBLE_DEVICE_ID_CHECK_FAIL = 0x18,
	PEBBLE_DEVICE_OK = 0x19,
	PEBBLE_DEVICE_COMPROMISED = 0x1A,
	PEBBLE_WRONG_POLICY_VERSION = 0x1B,
	PEBBLE_UNWRAP_DECAPSULATE_ERROR = 0x1C,
	PEBBLE_UNWRAP_APP_NAME_FAILED = 0x1D,
	PEBBLE_UNWRAP_FAILED_WRAP = 0x1E,
	PEBBLE_UNWRAP_FAILED_NO_KEY = 0x1F,
	PEBBLE_UNWRAP_FAILED = 0x20,
	PEBBLE_UNWRAP_VERSION_MISMATCH = 0x21,
	PEBBLE_UNWRAP_UNKNOWN = 0x22,
	PEBBLE_UNWRAP_LOAD_KEY_ERROR = 0x23,
	PEBBLE_UNWRAP_KEY_LENGTH_FAILED = 0x24,
	PEBBLE_ICCC_SAVE_STATUS_FAIL = 0x28,
	PEBBLE_JSON_FAIL = 0x29,
	PEBBLE_DEREGISTER_SHARED_BUF_FAIL = 0x2A,
	PEBBLE_REGISTER_SHARED_BUF_FAIL = 0x2B,
	PEBBLE_READ_EM_STATUS_FAIL = 0x30,
	PEBBLE_INVALID_CN_CHECK = 0x31,
	PEBBLE_NONCE_NOT_READY = 0x32,
	PEBBLE_GEN_NONCE_FAIL = 0x33,
	PEBBLE_WRAP_FAILED = 0x34,
	PEBBLE_ATN_NONCE_MISMATCH = 0x35,
	PEBBLE_GEN_CHECKSUM_FAIL = 0x36,
	PEBBLE_DRK_NOT_READY = 0x37,
	PEBBLE_CREATE_JWE_FAILED = 0x38,
	PEBBLE_CREATE_JWS_FAILED = 0x39,
	PEBBLE_EXTRACT_JWE_FAILED = 0x3A,
	PEBBLE_EXTRACT_JWS_FAILED = 0x3B,
	PEBBLE_SHARE_CHECKSUM_MISMATCH = 0x40,
	PEBBLE_INVALID_USER_INPUT_LEN = 0x41,
	PEBBLE_SSS_SPLIT_ERROR = 0x42,
	PEBBLE_SSS_JOIN_ERROR = 0x43,
} pebble_return_code_t;

/**
 * Normal world info
 */
typedef struct {
	uint8_t serial_number[NWS_INFO_LEN];
	uint8_t imei_0[NWS_INFO_LEN];
	uint8_t hash_imei[NWS_INFO_LEN];
	uint8_t mac_addr[NWS_INFO_LEN];
	uint8_t pebble_key[PEBBLE_DRK_MAX_BUF_LEN];
	uint32_t pebble_key_len;
	uint32_t is_wrapped_key;
} __attribute__ ((packed)) pebble_nwd_info_t;

/**
 * Service info
 */
typedef struct {
	uint8_t service_name[SERVICE_NAME_LEN];
	uint32_t service_request_len;
	uint8_t service_request_id[NWS_INFO_LEN];
} __attribute__ ((packed)) service_info_t;

/**
 * PEBBLE message
 */
typedef struct {
	uint32_t len;
	uint8_t data[JWS_LEN];
	uint32_t policy_value;
} __attribute__ ((packed)) tci_jws_message_t;

/**
 * load cert cmd
 */
typedef struct {
	BUF_T(PEBBLE_DRK_MAX_BUF_LEN)
	wrapped_key;
} __attribute__ ((packed)) load_cert_cmd_t;

/**
 * load cert resp
 */
typedef struct {
	cert_chain_t cert_chain;
} __attribute__ ((packed)) load_cert_resp_t;

/**
 * generate attestation challenge cmd
 */
typedef struct {
//	uint32_t len;  //nonce len is fixed by TA
} __attribute__ ((packed)) gen_atn_nonce_cmd_t;

/**
 * generate attestation challenge resp
 */
typedef struct {
	BUF_T(ATN_NONCE_MAX_LEN)
	nonce;
} __attribute__ ((packed)) gen_atn_nonce_resp_t;

/**
 * generate master key cmd.
 */
typedef struct {
	BUF_T(WRAPPING_AES_AAD_LEN)
	gcm_aad;
	cert_chain_t cert_chain; //certificate chain for AKS
} __attribute__ ((packed)) gen_master_key_cmd_t;

/**
 * generate master key resp
 */
typedef struct {
	BUF_T(SK_WRAP_KEY_LEN)
	master_key_wrapped_ta;
	BUF_T(SK_WRAP_KEY_LEN)
	secret_ct; // secret encrypted by AES GCM
	uint8_t gcm_iv[WRAPPING_AES_IV_LEN];
	uint8_t gcm_tag[WRAPPING_AES_TAG_LEN];
	BUF_T(WRAPPING_RSA_N_BYTES_LEN)
	wrapped_gcm_key;
} __attribute__ ((packed)) gen_master_key_resp_t;

/**
 * derive subkey cmd
 */
typedef struct {
	BUF_T(SK_WRAP_KEY_LEN)
	master_key_wrapped_ta;
	BUF_T(KEY_ALIAS_LEN)
	alias;
	BUF_T(WRAPPING_AES_AAD_LEN)
	gcm_aad;
	cert_chain_t cert_chain; //certificate chain for AKS
} __attribute__ ((packed)) derive_subkey_cmd_t;

/**
 * derive subkey resp
 */
typedef struct {
	BUF_T(SK_WRAP_KEY_LEN)
	secret_ct; // secret encrypted by AES GCM
	uint8_t gcm_iv[WRAPPING_AES_IV_LEN];
	uint8_t gcm_tag[WRAPPING_AES_TAG_LEN];
	BUF_T(WRAPPING_RSA_N_BYTES_LEN)
	wrapped_gcm_key;
} __attribute__ ((packed)) derive_subkey_resp_t;

#if 0 // split and join are for test only
/**
 * SSS split cmd
 */
typedef struct {
	uint32_t len;
	uint8_t data[SSS_KEY_MAX_LEN];
	uint8_t k;
	uint8_t n;
} __attribute__ ((packed)) split_cmd_t;

/**
 * SSS aplit resp
 */
typedef struct {
	uint32_t count;
	uint32_t len;
	uint8_t data0[SSS_SHARE_MAX_LEN];
	uint8_t data1[SSS_SHARE_MAX_LEN];
	uint8_t data2[SSS_SHARE_MAX_LEN];
} __attribute__ ((packed)) split_resp_t;

/**
 * SSS join cmd
 */
typedef struct {
	uint32_t count;
	uint32_t len;
	uint8_t data0[SSS_SHARE_MAX_LEN];
	uint8_t data1[SSS_SHARE_MAX_LEN];
	uint8_t data2[SSS_SHARE_MAX_LEN];
} __attribute__ ((packed)) join_cmd_t;

/**
 * SSS join resp
 */
typedef struct {
	uint32_t len;
	uint8_t data[SSS_KEY_MAX_LEN];
} __attribute__ ((packed)) join_resp_t;
#endif

/**
 * gen backup request cmd
 */
typedef struct {
	BUF_T(SK_WRAP_KEY_LEN)
	wrapped_master_key;
	cert_chain_t buddies_cert_chains[MAX_BUDDIES];
	uint32_t n; // ensure n <= MAX_BUDDIES
	uint32_t k;
	char metadata[KEY_METADATA_MAX_LEN];
} __attribute__ ((packed)) gen_backup_req_cmd_t;
/**
 * gen backup request resp
 */
typedef struct {
	char backup_reqs_jws[MAX_BUDDIES][JWE_JWS_MAX_SIZE];
	uint32_t n; // actual number of buddies/shares
} __attribute__ ((packed)) gen_backup_req_resp_t;

/**
 * gen_backup_ack cmd
 */
typedef struct {
	char backup_req_jws[JWE_JWS_MAX_SIZE];
	cert_chain_t buddy_cert_chain;
} __attribute__ ((packed)) gen_backup_ack_cmd_t;
/**
 * gen_backup_ack resp
 */
typedef struct {
	char backup_ack_jws[JWE_JWS_MAX_SIZE];
	BUF_T(SK_WRAP_KEY_LEN)
	wrapped_share;
	uint8_t n;
	uint8_t i;
	BUF_T(MAX_BUDDIES * CHECKSUM_LEN)
	checksums;
	char metadata[KEY_METADATA_MAX_LEN];
} __attribute__ ((packed)) gen_backup_ack_resp_t;

/**
 * gen_recovery_req cmd
 */
typedef struct {
	char key_alias[KEY_ALIAS_LEN];
} __attribute__ ((packed)) gen_recovery_req_cmd_t;
/**
 * gen_recovery_req resp
 */
typedef struct {
	char recovery_req_jws[JWE_JWS_MAX_SIZE];
} __attribute__ ((packed)) gen_recovery_req_resp_t;

/**
 * gen_recovery_ack cmd
 */
typedef struct {
	char recovery_req_jws[JWE_JWS_MAX_SIZE];
	BUF_T(SK_WRAP_KEY_LEN)
	wrapped_share;
	uint8_t n;
	uint8_t i;
	BUF_T(MAX_BUDDIES * CHECKSUM_LEN)
	checksums;
	char metadata[KEY_METADATA_MAX_LEN];
	cert_chain_t buddy_cert_chain;
} __attribute__ ((packed)) gen_recovery_ack_cmd_t;
/**
 * gen_recovery_ack resp
 */
typedef struct {
	char recovery_ack_jws[JWE_JWS_MAX_SIZE];
} __attribute__ ((packed)) gen_recovery_ack_resp_t;

/**
 * handle_recovery_ack cmd
 */
typedef struct {
	char recovery_ack_jws[JWE_JWS_MAX_SIZE];
	cert_chain_t buddy_cert_chain;
} __attribute__ ((packed)) handle_recovery_ack_cmd_t;
/**
 * handle_recovery_ack resp
 */
typedef struct {
	BUF_T(SK_WRAP_KEY_LEN)
	wrapped_share;
	uint8_t n;
	uint8_t i;
	BUF_T(MAX_BUDDIES * CHECKSUM_LEN)
	checksums;
	char metadata[KEY_METADATA_MAX_LEN];
} __attribute__ ((packed)) handle_recovery_ack_resp_t;

/**
 * reconstruct_key cmd
 */
typedef struct {
	uint8_t wrapped_shares[MAX_BUDDIES][SK_WRAP_KEY_LEN]; // does not contain metadata
	uint8_t k; // k shares are stored in the top k rows of wrapped_shares
	uint32_t wshare_len; // actual len of all wrapped shares, assuming they have the same len
	gen_master_key_cmd_t master_key_req;
} __attribute__ ((packed)) reconstruct_key_cmd_t;
/**
 * reconstruct_key resp
 */
typedef struct {
	gen_master_key_resp_t master_key_resp;
} __attribute__ ((packed)) reconstruct_key_resp_t;

typedef union {
	load_cert_cmd_t load_cert_cmd;
	load_cert_resp_t load_cert_resp;

	gen_atn_nonce_cmd_t gen_atn_nonce_cmd;
	gen_atn_nonce_resp_t gen_atn_nonce_resp;

	gen_master_key_cmd_t gen_master_key_cmd;
	gen_master_key_resp_t gen_master_key_resp;

	derive_subkey_cmd_t derive_subkey_cmd;
	derive_subkey_resp_t derive_subkey_resp;

//	split_cmd_t split_cmd;  //TODO: for testing, need to disable
//	split_resp_t split_resp; //TODO: for testing, need to disable
//
//	join_cmd_t join_cmd; //TODO: for testing, need to disable
//	join_resp_t join_resp; //TODO: for testing, need to disable

	gen_backup_req_cmd_t gen_backup_req_cmd;
	gen_backup_req_resp_t gen_backup_req_resp;

	gen_backup_ack_cmd_t gen_backup_ack_cmd;
	gen_backup_ack_resp_t gen_backup_ack_resp;

	gen_recovery_req_cmd_t gen_recovery_req_cmd;
	gen_recovery_req_resp_t gen_recovery_req_resp;

	gen_recovery_ack_cmd_t gen_recovery_ack_cmd;
	gen_recovery_ack_resp_t gen_recovery_ack_resp;

	handle_recovery_ack_cmd_t handle_recovery_ack_cmd;
	handle_recovery_ack_resp_t handle_recovery_ack_resp;

	reconstruct_key_cmd_t reconstruct_key_cmd;
	reconstruct_key_resp_t reconstruct_key_resp;

//  tci_jws_message_t jws_message;
} __attribute__ ((packed)) payload_t;

/**
 * TCI message
 */
typedef struct {
	tz_msg_header_t header;
	payload_t payload;
} __attribute__ ((packed)) tci_message_t;

#endif /* _TZ_PEBBLE_INTERFACE_H_ */