typedef unsigned char u8;
typedef unsigned int u32;
typedef unsigned long long u64;
 
#include <stdio.h>
#include <unistd.h>
#include <sys/socket.h>
#include <fcntl.h>
#include <sys/ioctl.h>
#include <linux/if_alg.h>
#include <linux/socket.h>
#include <linux/types.h>
#include "xlnx_puf.h"
#include <string.h>
#include <errno.h>
#include <stdbool.h>
#include <sys/types.h>
#include <unistd.h>
#include <stdlib.h>
 
 
#define MAX_SYNDROME_DATA_LEN  (PUF_MAX_SYNDROME_DATA_LEN_IN_WORDS * sizeof(u32))
#define MAX_EFUSE_TRIM_SYN_DATA_LEN  (PUF_EFUSE_TRIM_SYN_DATA_IN_WORDS * sizeof(u32))
#define MAX_PUF_ID_LEN (PUF_ID_LEN_IN_WORDS * sizeof(u32))
 
#define PUF_SHUTTER_VALUE 0x81000100U
 
/* Xilinx error codes */
#define RD_FAILED     1026
#define WR_FAILED     1027
  
#define SYS_PATH    "/sys/bus/nvmem/devices/xilinx-secure-config0/nvmem"
 
#ifndef SOL_ALG
#define SOL_ALG 279
#endif
 
#define KEY_SIZE        32
#define IV_SIZE         12
#define AAD_SIZE    16
#define GCM_TAG_SIZE    16
#define DATA_SIZE   32
#define PUF_WORD_LENGTH 0x4
#define PUF_RED_KEY_LEN_IN_BYTES 32
#define PUF_RED_KEY_LEN_IN_BITS PUF_RED_KEY_LEN_IN_BYTES*8
#define HASH_STRING_LEN_IN_BYTES 65
#define BLACK_IV_STRING_LEN_IN_BYTES 25
#define PUF_DATA_LENGTH 129
#define PUF_BLACK_KEY_STRING_LEN_IN_BYTES 64
#define PUF_DATA_LENGTH_IN_BYTES    512
#define PUF_DATA_STR_LEN_IN_BYTES   1016
 
#define VERSAL_AES_PUF_kEY 11 /* Refer enum in driver for more details */
 struct pufdata puf_data;
  char str[1016] = {0};
int PufRegis(void) {
    int fd;
    int err;
    struct puf_usrparams params;
    
    ssize_t size;
     
    // Open the PUF device file
    fd = open("/dev/xpuf", O_RDWR);
    if (fd < 0) {
        perror("Failed to open PUF device /dev/xlnx-puf");
        err = 1;
        goto END;
    }
  
    // Set the parameters for PUF registration
    params.pufoperation = 0U;  // PUF registration
    params.globalvarfilter = 1U;
    params.readoption = PUF_READ_FROM_RAM;
    params.shuttervalue = PUF_SHUTTER_VALUE;
    params.pufdataaddr = (u64)&puf_data.pufhd;
    params.trimsyndataaddr = (u64)puf_data.efusesyndata;
     
    // Perform PUF registration
    if (ioctl(fd, PUF_REGISTRATION, &params) < 0) {
        perror("PUF registration failed \n");
        err = 1;
        goto END;
    }
     // Print the PUF ID
    printf("1 : PUF ID registartion:\n ");
    for (int i = 0; i < PUF_ID_LEN_IN_WORDS * 4U; i++) {
        printf("%02x ", *((u8*)puf_data.pufid + i));
    }
    printf("\n");
 
    // Print the syndrome data
    printf("Syndrome Data:\n");
    for (int i = 0; i < PUF_MAX_SYNDROME_DATA_LEN_IN_WORDS; i++) {
        printf("%08x ", puf_data.pufhd.syndata[i]);
    }
    printf("\n");
     
    printf("trimmed Syndrome Data:\n");
    for (int i = 0; i < PUF_EFUSE_TRIM_SYN_DATA_IN_WORDS; i++) {
        printf("%08x ", puf_data.efusesyndata[i]);
    }
    printf("\n");
     
    printf("PUF id:\n");
    for (int i = 0; i < PUF_ID_LEN_IN_WORDS; i++) {
        printf("%08x ", puf_data.pufid[i]);
    }
    printf("\n");
     
    // Print the auxiliary data
    printf("Auxiliary Data: %08x\n", puf_data.pufhd.aux);
    printf("Chash Data: %08x\n", puf_data.pufhd.chash);
 
    close(fd);
    err = 0;
    printf("PUF registration completed successfully.\n");
END:
    return err;
}
 
void FormatAesKey(const u8* Key, u8* FormattedKey, u32 KeyLen)
{
int Index;
int Words = KeyLen / PUF_WORD_LENGTH;
u32 WordIndex = Words / 2U;
    u32* InputKey = (u32*)Key;
    u32* OutputKey  = (u32*)FormattedKey;
 
for(Index = 0U; Index < Words; Index++)
    {
        OutputKey[Index] = InputKey[WordIndex];
        WordIndex++;
        /*
         * AES word size = 128 bits
         * So to change the endianness, code should swap lower 64bits
         * with upper 64 bits
         * 64 bits = 8 bytes
         */
        WordIndex = WordIndex % 8U;
    }
}
 
void ReverseData (const u8 *OrgDataPtr, u8* SwapPtr, u32 Len){
u32 Index = 0U;
    u32 ReverseIndex = Len - 1U;
 
    for(Index = 0U; Index < Len; Index++)
    {
        SwapPtr[Index] = OrgDataPtr[ReverseIndex];
        ReverseIndex--;
    }
}
 
int main(void)
{
    int len = DATA_SIZE + GCM_TAG_SIZE;
    int opfd;
    int nvmfd;
    ssize_t size;
    int tfmfd;
    unsigned int keytype;
    unsigned int envdis;
    u_int64_t offset;
    struct sockaddr_alg sa = {
        .salg_family = AF_ALG,
        .salg_type = "aead",
        .salg_name = "gcm(aes)"
    };
    struct msghdr msg = {};
    struct cmsghdr *cmsg;
    char cbuf[CMSG_SPACE(4) + CMSG_SPACE(1024) + CMSG_SPACE(1024)] = {0};
    char buf[len];
    int length = strlen(str);
    struct af_alg_iv *iv;
    struct iovec iov;
    int i;
    u32 PufData[PUF_DATA_LENGTH];
    unsigned char blackkey[PUF_BLACK_KEY_STRING_LEN_IN_BYTES] = {0};
    unsigned char blkkey[PUF_BLACK_KEY_STRING_LEN_IN_BYTES]={0};
    char PufDataresult[PUF_DATA_LENGTH_IN_BYTES]={0};
    char PufDataInString[PUF_DATA_STR_LEN_IN_BYTES] = {0};
    char hash[HASH_STRING_LEN_IN_BYTES]="5CA45B3D17F4F809549796A9ECE8CC23365AFAC4BCA636B11F08ABE2EF54C36F";
    char blackiv[BLACK_IV_STRING_LEN_IN_BYTES] = "D2450E07EA5DE0426C0FA133";
    u8 FormattedBlackKey[PUF_RED_KEY_LEN_IN_BITS] = {0};
    u8 FlashBlackKey[PUF_RED_KEY_LEN_IN_BITS]  = {0};
    unsigned char temp[len];
    u8 temp1[PUF_RED_KEY_LEN_IN_BITS];
     
     
    __u8 dummykey[32U] = {0U};
 
    __u8 key[] = {
        0xF8, 0x78, 0xB8, 0x38, 0xD8, 0x58, 0x98, 0x18, 0xE8, 0x68, 0xA8, 0x28, 0xC8, 0x48,
        0x88, 0x08, 0xF0, 0x70, 0xB0, 0x30, 0xD0, 0x50, 0x90, 0x10, 0xE0, 0x60, 0xA0, 0x20,
        0xC0, 0x40, 0x80, 0x00
    };
    __u8 key_type[] = {VERSAL_AES_PUF_kEY};
 
    __u8 usr_iv[] = {
        0xD2, 0x45, 0x0E, 0x07, 0xEA, 0x5D, 0xE0, 0x42, 0x6C, 0x0F, 0xA1, 0x33
 
    };
    __u8 input[DATA_SIZE] = {
         
        /* Plain text */
        0x12, 0x34, 0x56, 0x78, 0x08, 0xF0, 0x70, 0xB0, 0x30, 0xD0, 0x50, 0x90, 0x10, 0xE0,
        0x60, 0xA0, 0x20, 0xC0, 0x40, 0x80, 0x00, 0xA5, 0xDE, 0x08, 0xD8, 0x58, 0x98, 0xA5,
        0xA5, 0xFE, 0xDC, 0xA1
    };
 
    /* Do registration */
    if (PufRegis() > 0) {
        perror("PUF registartion failed\n");
        return 1;
    }
 
    /* Do enc/dec using PUF kek */
    tfmfd = socket(AF_ALG, SOCK_SEQPACKET, 0);
    bind(tfmfd, (struct sockaddr *)&sa, sizeof(sa));
    setsockopt(tfmfd, SOL_ALG, ALG_SET_KEY, key_type, 1);
    setsockopt(tfmfd, SOL_ALG, ALG_SET_KEY, dummykey, 32);
    setsockopt(tfmfd, SOL_ALG, ALG_SET_AEAD_AUTHSIZE, NULL, 16);
    //setsockopt(tfmfd, SOL_ALG, ALG_SET_AEAD_ASSOCLEN, aad, AAD_SIZE);
    opfd = accept(tfmfd, NULL, 0);
    msg.msg_control = cbuf;
    msg.msg_controllen = sizeof(cbuf);
    cmsg = CMSG_FIRSTHDR(&msg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_OP;
    cmsg->cmsg_len = CMSG_LEN(4);
    *(__u32 *)CMSG_DATA(cmsg) = ALG_OP_ENCRYPT;
    cmsg = CMSG_NXTHDR(&msg, cmsg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_IV;
    cmsg->cmsg_len = CMSG_LEN(1024);
 
    iv = (void *)CMSG_DATA(cmsg);
    iv->ivlen = IV_SIZE;
    memcpy(iv->iv, usr_iv, IV_SIZE);
 
    /* AAD may not be required in case of KEK enc/dec use case */
    cmsg = CMSG_NXTHDR(&msg, cmsg);
    cmsg->cmsg_level = SOL_ALG;
    cmsg->cmsg_type = ALG_SET_AEAD_ASSOCLEN;
    cmsg->cmsg_len = CMSG_LEN(1024);
    iv = (void *)CMSG_DATA(cmsg);
    //iv->ivlen = AAD_SIZE;
 
    iov.iov_base = input;
    iov.iov_len = len;
 
    msg.msg_iov = &iov;
    msg.msg_iovlen = 1;
 
    sendmsg(opfd, &msg, 0);
    read(opfd, buf, len);
     
    printf("Data Out:\r\n");
    for (i = 0; i < len - GCM_TAG_SIZE; i++) {
        printf("%02x", (unsigned char)buf[i]);
    }
    printf("\n");
 
    printf("GCM TAG:\r\n");
    for (i = len - GCM_TAG_SIZE; i < len; i++) {
        printf("%02x", (unsigned char)buf[i]);
    }
 
    printf("\n");
 
    printf("Input:\r\n");
    for (i = 0; i < len - GCM_TAG_SIZE; i++) {
        printf("%02x", (unsigned char)input[i]);
    }
    close(opfd);
    close(tfmfd);
     
     
    nvmfd = open(SYS_PATH, O_RDWR);
    printf("nvmfd = %x \n\r",nvmfd);
    if(nvmfd <= 0) {
        printf("Opening SYS FS NVMEM file is failed %d \n", errno);
        perror(open);
        return errno;
    }
 
    /* programming black iv */
 
    size = pwrite(nvmfd,blackiv,2*IV_SIZE,0x301D0);
    if (size == 2*IV_SIZE) {
        printf("Black IV PROGRAMMED \n\r");
    }
    else {
        printf("Black IV is not PROGRAMMED \n\r");
        return WR_FAILED;
    }
 
 
    size = pwrite(nvmfd,hash,2*DATA_SIZE,0x30100);
    if (size == 2*DATA_SIZE) {
        printf("PPK HASH IS PROGRAMMED\n\r");
    }
    else {
        printf("PPK HASH is not PROGRAMMED \n\r");
        return WR_FAILED;
    }
     
    PufData[0] = puf_data.pufhd.chash;
    PufData[1] = puf_data.pufhd.aux;
     
     
    for(int i=2;i<129;i++) {
        PufData[i]= puf_data.efusesyndata[i-2];
    }
     
    for(int i=0;i<129;i++){
    sprintf(PufDataresult,"%08x",PufData[i]);
    strcat(PufDataInString,PufDataresult);
    }
 
     
     
    length = strlen(PufDataInString);
    size = pwrite(nvmfd,PufDataInString,length,0x3FFFF);
 
    if (size == length) {
        printf("PUF data is PROGRAMMED \n\r");
    }
    else {
        printf("PUF data is not PROGRAMMED \n\r");
        return WR_FAILED;
    }
     
     
     
    /*buf is an array where encrypted data(encrypted redkey) is stored*/
    FormatAesKey(buf,FormattedBlackKey,PUF_RED_KEY_LEN_IN_BYTES);
    ReverseData(FormattedBlackKey, FlashBlackKey, PUF_RED_KEY_LEN_IN_BYTES);
     
    for(int i=0;i<DATA_SIZE;i++){
    sprintf(blackkey,"%02x",(unsigned char)FlashBlackKey[i]);
    strcat(blkkey,blackkey);
    }
     
    /*Key type can be choosen accordingly AES KEY = 1, USER_KEY_0 = 2, USER_KEY_1 = 4*/
    keytype = 1;
    envdis = 1;
    offset = (1 << 16)| (envdis<<17) | keytype;
    size = pwrite(nvmfd, blkkey, 2*DATA_SIZE, offset);
     
    if (size == 2*DATA_SIZE) {
        printf("Black key is programmed \n\r");
    }
    else {
        printf("blackkey programming failed \n\r");
        return WR_FAILED;
    }
     
     
    close(nvmfd);
     
     
 
 
    return 0;
}