/*
 * The contents of this file are subject to the terms
 * of the Common Development and Distribution License
 * (the "License").  You may not use this file except
 * in compliance with the License.
 *
 * You can obtain a copy of the license at
 * http://www.opensource.org/licenses/cddl1.php
 * See the License for the specific language governing
 * permissions and limitations under the License.
 *
 * When distributing Covered Code, include this CDDL
 * HEADER in each file and include the License file at
 * http://www.opensource.org/licenses/cddl1.php.  If 
 * applicable, add the following below this CDDL HEADER, 
 * with the fields enclosed by brackets "[]" replaced 
 * with your own identifying information: 
 * Portions Copyright [yyyy]
 * [name of copyright owner]
 */ 

/*
 * $(@)main.c $Revision: 1.3 $ $Date: 2006/09/29 18:17:39 $
 * 
 * Copyright 2006 Sun Microsystems, Inc. All Rights Reserved.
 */
/*
 * Copyright (c) 2005 by Sun Microsystems, Inc.
 * Functions to perform TS parsing

 * Created:   July 29, 2005
 * Author:    Yongfang Liang
 *---------------------------------------------------------------*/

#include <stdio.h>

//#include "trickplay/trickplay.h"
#include "cas/cas_lib.h"
#include "ca_crypto/crypto.h"
#define HAVE_ENCRYPTION_LIB
//#define _DEBUG

#include "ecmfile.h"
#include <libgen.h>
#include <stdio.h> 
#include <stdlib.h>
#include <sys/types.h>
#include <fcntl.h>
#include <unistd.h>
#include "cobjs/Macros.h"
#include <sys/mman.h>

#include "cobjs/ArrayOf.h"
#include "cobjs/HashTable.h"
#include "cobjs/Macros.h"
#include "cobjs/RetCode.h"
#include "cobjs/Types.h"
#include "cobjs/libcobjsVersion.h"

#include "mmp/Mmp.h"
#include "mmp/Mpeg1Pes.h"
#include "mmp/Mpeg1Sys.h"
#include "mmp/Mpeg2.h"
#include "mmp/Mpeg2Pat.h"
#include "mmp/Mpeg2Pes.h"
#include "mmp/Mpeg2Pmt.h"
#include "mmp/Mpeg2Ps.h"
#include "mmp/Mpeg2Psi.h"
#include "mmp/Mpeg2Ts.h"
#include "mmp/Pipe.h"
#include "mmp/libmmpVersion.h"
#include "mmp/Mpeg2VideoEs.h"
#include "mmp/Mpeg1AudioEs.h"

#define CA_SYS_ID  0x222
#define FREQUENCY 5
/* AES key size 128 bits */
#define KEY_LENGTH 16 
/* AES use block size 16bytes */
#define BLOCK_SIZE 16


typedef struct 
{
    uint8_t  key_len;
    uint8_t  iv_len;
    uint8_t  key[KEY_LENGTH];
    uint8_t  iv[BLOCK_SIZE];
} cas_ecm_payload, *cas_ecm_payload_p;

typedef struct 
{
    uint8_t  aes_mode;
    uint8_t  key_len;
    uint8_t  iv_len;
    uint8_t  key[KEY_LENGTH];
    uint8_t  iv[BLOCK_SIZE];
} cas_emm_payload, *cas_emm_payload_p;


typedef struct 
{
	cas_ecm_payload  *p_ecm;
	cas_emm_payload  *p_emm;
	/* input ecm file name,only used for generating encrypted Trickplay stream */
	char in_ecm_file_name[256]; 
	/* output ecm recording the ECM offset in the encrypted file 
	   used when generating encrypted Trickplay stream */ 
	char out_ecm_file_name[256];
	char emm_file_name[256];
	int  ecm_no;
	int  ecm_no_inserting;
    ca_sys_p p_ecm_sys;
	short ecm_id;
	int b_reverse;

} sys_info_s, *sys_info_p;

#define	MAP_CHUNK	(1024*1024)
static void  OutputCallBack (void *token, uint8_t *buf, int len);

static void  ECMCallBack (void *token, uint8_t *buf, int len, uint8_t *out_buf, int *out_len, long long  position);
static void  ECMIOnlyCallBack (void *token, uint8_t *buf, int len, uint8_t *out_buf, int *out_len, long long  position);

static void  ScrambleCallBack (void *token, uint8_t *key_buf, int key_len, 
	                       uint8_t *pkt, uint8_t *payload_buf, int payload_len,
	                       long long position);

static void  EMMOutputCallBack (void *token, uint8_t *ts_buf, int ts_len);
static void EMMAsymEncrypt(void *token, uint8_t *emm_buf, int emm_len, 
                        uint8_t *out_buf, int *out_len, long long  position);
static void TSCallBack(void *token, long long position);

/* put the file dat into a MMP object */
static long long file2mmp(const char *file, ca_sys_p p_sys)
{
    long                pagesize = sysconf(_SC_PAGESIZE);
    off64_t             offset = 0;
    long                mapChunk;
    int                 fd;
    struct stat64       st;
    caddr_t             buf;
    long long           retVal;
    int		retCode;

    if ((fd = open64(file, O_RDONLY)) <= 0 || fstat64(fd, &st) < 0) {
	perror(file);
	retVal = -1;
	goto done;
    }
    retVal = st.st_size;
    (void) printf("File size:     %lld\n", st.st_size);

    mapChunk = ROUND_UP(MAP_CHUNK, pagesize);
    while (offset < st.st_size) {
	buf = mmap64((caddr_t) 0, mapChunk, PROT_READ, MAP_PRIVATE, fd, offset);
	if (buf == (caddr_t) - 1) {
	    perror("mmap");
	    retVal = -1;
	    goto done;
	}
	if (madvise(buf, mapChunk, MADV_SEQUENTIAL) < 0) {
	    perror("madvise");
	    retVal = -1;
	    goto done;
	}
	retCode = casPutData(p_sys, (uint8_t *)buf,
		(size_t) MIN(mapChunk, st.st_size - offset));
	if (munmap(buf, mapChunk) < 0) {
	    perror("munmap");
	    retVal = -1;
	    goto done;
	}
	if (retCode != 0) {
	    break;
	}
	offset += mapChunk;
    }
done:
    if (fd >= 0) {
	(void) close(fd);
    }
    return retVal;
}

static void  EMMOutputCallBack (void *token, uint8_t *ts_buf, int ts_len)
{
    /* generate the EMM key file */
    FILE *emmFile = NULL;
    char *filename = (char *)token;
    emmFile = fopen(filename,"wb");
    fwrite(ts_buf, 1, ts_len, emmFile);
    fclose(emmFile);
    return;
}


static void EMMWritePayloadCallBack(void *token, uint8_t *emm_buf, int emm_len, 
                            uint8_t *out_buf, int *out_len, long long  position)
{
    FILE *emmFile = NULL;
    char *filename = (char *)token;
    emmFile = fopen(filename,"wb");
    fwrite(emm_buf, 1, emm_len, emmFile);
    fclose(emmFile);

    /* 
     * Hacky!
     * function expects something in the out_buf hence copy contents into
     * out_buf 
     */
    memcpy(out_buf, emm_buf, emm_len);
    *out_len = emm_len;
    return;
 
}
/* generate the EMM data file */

static void createEMM(cas_emm_payload *p_emm, uint8_t aes_mode, char *filename)
{
    ca_sys_p            p_sys, p_emm_sys;
    int i = 0;
#ifndef HAVE_ENCRYPTION_LIB
    unsigned char access_key[KEY_LENGTH+1]  = "2222222222222222";
    unsigned char access_iv[BLOCK_SIZE+1]   = "0123456789012345";
    memcpy(p_emm->iv, access_iv, BLOCK_SIZE);
    memcpy(p_emm->key, access_key, KEY_LENGTH);
#else
    generate_key(p_emm->iv, BLOCK_SIZE);
    generate_key(p_emm->key, KEY_LENGTH);
#endif

    p_emm->key_len = KEY_LENGTH ;
    p_emm->iv_len  = BLOCK_SIZE;
    p_emm->aes_mode = aes_mode;
    
#ifdef _DEBUG
    printf ("key: ");
    for (i = 0; i<KEY_LENGTH; i++)
    {
        printf("%#x ", p_emm->key[i]);
    }
    printf("\n");
    printf ("iv : ");
    for (i = 0; i<BLOCK_SIZE; i++)
    {
        printf("%#x ", p_emm->iv[i]);
    }
    printf("\n");
#endif
    
    p_emm_sys = createCAS(CA_SYS_ID,FREQUENCY);
    /* we dont want to write the whole TS packet. Just the payload. */
    // casSetOutputCallback(p_emm_sys, EMMOutputCallBack, filename);

    casSetEMSCallBack(p_emm_sys, EMMWritePayloadCallBack, filename);
    generateCAMessage(p_emm_sys, 0, 600, (uint8_t *)p_emm, sizeof(*p_emm));
    freeCAS(p_emm_sys);
    return;
	
}

void GetEncryptedBitstream(ca_sys_p  p_sys, sys_info_s *p_sys_info, char *infilename, FILE *outfile, int ecm_after_pmt,
	      casScrambleCallBack f_scramble,
	      casEMSCallBack      f_ems
	      )
{
    long long           fileSize;
    int             retCode;
    
    casECMafterPMT(p_sys, ecm_after_pmt);
    
    casSetOutputCallback(p_sys, OutputCallBack, outfile);
    casSetEMSCallBack(p_sys, f_ems, p_sys_info);
    casSetScrambleCallBack(p_sys, f_scramble, p_sys_info);
    
    addKey2CAS(p_sys, FOR_ANY_PROGRAM, ENCRYPT_ALL, 
    	       (uint8_t *)p_sys_info->p_ecm, sizeof(cas_ecm_payload));
//    addKey2CAS(p_sys, FOR_ANY_PROGRAM, ENCRYPT_VIDEO_ONLY, (uint8_t *)&ecm, sizeof( cas_ecm_payload));
    
    /* read file, parse it, and then generate the new file */
    printf("Start ECM insertion...\n");
    fileSize = file2mmp(infilename, p_sys);
    retCode = casEOF(p_sys);
    printf("Done!\n");
    return;
}

static void  OutputCallBack (void *token, uint8_t *ts_buf, int ts_len)
{
	FILE *outfile = (FILE *)token;
	fwrite(ts_buf, 1, ts_len, outfile);
	return;
}

/* position -> the postion of ECM in the newfile */
/* key_buf: the original payload added by addKey2CAS */
/* out_buf is the buffer being inserted into the bitstream */
/* this function is called each time a ecm is generated */
static void  ECMCallBack (void *token, uint8_t *key_buf, int key_len, uint8_t *out_buf, int *out_len, long long  position)
{
    sys_info_s *p_sys_info = (sys_info_s *)token;
    FILE *ecm_file = NULL;
    cas_ecm_payload_p p_ecm = p_sys_info->p_ecm;
    cas_emm_payload_p p_emm = p_sys_info->p_emm;
    
    int encrypted_data_size = 0;
    uint8_t *p_ecm_buf = (uint8_t *)(p_ecm);
    
#ifndef HAVE_ENCRYPTION_LIB
    unsigned char access_key[KEY_LENGTH+1]  = "2222222222222222";
    unsigned char access_iv[BLOCK_SIZE+1]   = "0123456789012345";
    memcpy(p_ecm->iv, access_iv, BLOCK_SIZE);
    memcpy(p_ecm->key, access_key, KEY_LENGTH);
#else
    generate_key(p_ecm->iv, BLOCK_SIZE);
    generate_key(p_ecm->key, KEY_LENGTH);
#endif

    /* when everytime inserting the ECM, the key and IV are changed */
    p_ecm->key_len = KEY_LENGTH ;
    p_ecm->iv_len  = BLOCK_SIZE;
    


    /* save the ECM payload for later used in encrypting the trickplay stream */
    ecm_file = fopen(p_sys_info->out_ecm_file_name, "r+b");
    if(ecm_file)
    {
        ECMfile_info *info = GetECMFileInfo(ecm_file);
        /* save the ECM, not yet encrypted by EMM! */
        WriteOneECMItem(ecm_file, info, p_sys_info->ecm_no,
                        p_ecm_buf, sizeof(*p_ecm), position);
    	p_sys_info->ecm_no ++;
        fclose(ecm_file);
        CleanECMInfo(info);
    }

#ifdef HAVE_ENCRYPTION_LIB
    encrypted_data_size =  encryptBufferData(p_ecm_buf, sizeof(cas_ecm_payload), out_buf, *out_len, 
    	              p_emm->key, p_emm->key_len, p_emm->iv, p_emm->iv_len, (AES_MODE)p_emm->aes_mode);
    *out_len = encrypted_data_size;
#else
    memcpy(out_buf, p_ecm_buf, *out_len);
    *out_len = sizeof(cas_ecm_payload);
    
#endif

    return;
}

/* position -> the postion of ECM in the newfile */
/* key_buf: the original payload added by addKey2CAS */
/* out_buf is the buffer being inserted into the bitstream */
/* this function is called each time a ecm is generated */
static void  ECMIOnlyCallBack (void *token, uint8_t *key_buf, int key_len, uint8_t *out_buf, int *out_len, long long  position)
{
    sys_info_s *p_sys_info = (sys_info_s *)token;
    FILE *ecm_file = NULL;
    cas_ecm_payload_p p_ecm = p_sys_info->p_ecm;
    cas_emm_payload_p p_emm = p_sys_info->p_emm;
    
    int encrypted_data_size = 0;

    memcpy(p_ecm, key_buf, key_len);

#ifdef _DEBUG
    ecm_file = fopen(p_sys_info->out_ecm_file_name, "r+b");
    if(ecm_file)
    {
        ECMfile_info *info = GetECMFileInfo(ecm_file);
        /* save the ECM, not yet encrypted by EMM! */
        WriteOneECMItem(ecm_file, info, p_sys_info->ecm_no,
                        key_buf, key_len, position);
    	p_sys_info->ecm_no ++;
        fclose(ecm_file);
        CleanECMInfo(info);
    }
#endif

#ifdef HAVE_ENCRYPTION_LIB
    encrypted_data_size =  encryptBufferData(key_buf, key_len, out_buf, *out_len, 
    	              p_emm->key, p_emm->key_len, p_emm->iv, p_emm->iv_len, (AES_MODE)p_emm->aes_mode);
    *out_len = encrypted_data_size;
#endif
    return;
}

static void counterAddUInt32( uint8_t *op1, uint32_t op_size, uint32_t op2)
{
    uint32_t len;
    uint8_t  op2_8; 
    
    len = op_size - 1;
    op2_8 = op2 & 0xFF;
    while ( len >= 0 )
    {
        op1[len] = op1[len] + op2_8;
    	if ( op1[len] )
    		break;
    	len--;
    }

    len = op_size - 2;
    op2_8 = (op2 >> 8 )& 0xFF;
    while ( len >= 0 )
    {
        op1[len] = op1[len] + op2_8;
    	if ( op1[len] )
    		break;
    	len--;
    }

    len = op_size - 3;
    op2_8 = (op2 >> 16 ) & 0xFF;
    while ( len >= 0 )
    {
        op1[len] = op1[len] + op2_8;
    	if ( op1[len] )
    		break;
    	len--;
    }

    len = op_size - 4;
    op2_8 = (op2 >> 24 ) & 0xFF;
    while ( len >= 0 )
    {
        op1[len] = op1[len] + op2_8;
    	if ( op1[len] )
    		break;
    	len--;
    }
    
    return;  	
}

static void counterInc( uint8_t *op1, uint32_t op_size)
{
    uint32_t len;
    uint8_t  op2_8; 
    
    len = op_size - 1;
    while ( len >= 0 )
    {
    	if ( ++op1[len] )
    		break;
    	len--;
    }
   
    return;  	
}

static void TSCallBack(void *token, long long position)
{
    sys_info_s *p_sys_info = (sys_info_s *)token;
    cas_ecm_payload_p p_ecm = p_sys_info->p_ecm;
    FILE *ecm_file = NULL;
    
    /* inserting the ECM at proper positions */
    ecm_file = fopen(p_sys_info->in_ecm_file_name, "rb");
    if (ecm_file == NULL)
    {
        printf("Error: could not open ECM file, no ECM inserted!\n");
    	return;
    }
    else 
    {
        ECMfile_info *info;
        int ecm_index = 0;
        info = GetECMFileInfo(ecm_file);
        ecm_index = p_sys_info->b_reverse ? info->header.no_of_ECMs-1-p_sys_info->ecm_no_inserting : p_sys_info->ecm_no_inserting;
        
        ReWrOneECMItem(ecm_file, info, ecm_index, TRUE);
    	while (info->ecmItem.insert_offset == -1