/* vim:set sw=4 ts=8 fileencoding=cp1251::Кодировка:WINDOWS-1251[АБЁЪЯабёъя] */
#ifdef _WIN32
    #pragma setlocale("rus")
#endif
/*
 * Copyright(C) 2013-2014 Проект ИОК
 *
 * Этот файл содержит информацию, являющуюся
 * собственностью компании Крипто Про.
 *
 * Любая часть этого файла не может быть скопирована,
 * исправлена, переведена на другие языки,
 * локализована или модифицирована любым способом,
 * откомпилирована, передана по сети с или на
 * любую компьютерную систему без предварительного
 * заключения соглашения с компанией Крипто Про.
 *
 * Программный код, содержащийся в этом файле, предназначен
 * исключительно для целей обучения и не может быть использован
 * для защиты информации.
 *
 * Компания Крипто-Про не несет никакой
 * ответственности за функционирование этого кода.
 */

/*!
 * \version $Revision: 89611 $
 * \date $Date: 2013-05-29 23:43:11 +0400 (СЃСЂ, 29 РјР°Р№ 2013) $
 * \author $Author: lse $
 *
 * \brief Пример использования реализации TLS уровня ядра Linux
 */
#ifdef CSP_DRIVER 
typedef __WCHAR_TYPE__ wchar_t;
#endif
#include"KTLS.h"
#include"wincspc.h"
#include "WCKernelPart.h"
#ifndef CSP_DRIVER
#include"WCMemSTDC.h"
#endif
#include"drtktls_io.h"

#if defined(KRN_EMUL)
	// Userspace Example with libcpcdrv_emul.a and libktls_emul.a
    #include <sys/socket.h>
    #include <errno.h>
    #include <string.h>
    #include <stdarg.h>
    #include <stdio.h>
    #define NATT 10

    static ssize_t
    emulate_recv(int socket, void *buffer, size_t length, int flags)
    {
	ssize_t ret;
	int natt = 0;
	do 
	{
	    struct timespec tvi,tvo;
	    ret=recv(socket, buffer, length, flags);
	    (ret==-1)?-errno:ret;
	    if ( ret != -EINTR && ret != -EAGAIN )
		break;
	    if ( ++natt == NATT )
	    {
	    	fprintf(stderr,"drtktls_emul:fail_recv\n");
		break;
	    }

	    tvi.tv_sec=0;
	    tvi.tv_nsec=(1000<<natt);
	    nanosleep(&tvi,&tvo);
	}
	while (1);	
	return ret;
    }
#else
#ifdef CSP_DRIVER
    #include <linux/kernel.h>
    #include <linux/string.h>
    #include <linux/errno.h>
    #define mark_alloced()
    #define check_alloced()
#endif    
    ssize_t
    emulate_recv(int socket, 
		void *buffer, 	// buffer in kernel space
		size_t length, 
		int flags);
#endif

DWORD totalAlloced=0;

void * CPCAlloc(LPCPC_CONFIG CSPConfig, size_t dwSize, DWORD dwMemPoolId)
{
    unsigned long dwThreadId;
    void * res=NULL;
    DWORD * pdw;
    if (CSPConfig == NULL)
	return NULL;
    if (CSPConfig->pArena == NULL)
	return NULL;
    if (CSPConfig->logConfig.get_thread_id == NULL)
	dwThreadId=0;
    else
	dwThreadId=CSPConfig->logConfig.get_thread_id();
    if (CSPConfig->pArena->pAllocMemory(
			CSPConfig->pArena,
			dwSize+16,
			dwMemPoolId,
			dwThreadId,
			&res)!=S_OK)
	return NULL;
    memset(res,0,dwSize+16);
    pdw=(DWORD *)res;
    *pdw=dwSize;
    totalAlloced+=dwSize;
    return pdw+4;
}

/* Освободить память с помощью менеджера памяти, заданного пользователем */
void CPCFree(LPCPC_CONFIG CSPConfig, void * ptr, DWORD dwMemPoolId)
{
    if(ptr)
    {
	DWORD * pdw=(DWORD *)ptr-4;
	totalAlloced-=*pdw;
	CSPConfig->pArena->pFreeMemory(CSPConfig->pArena,pdw,dwMemPoolId);
    }
}

#define PRINTF_SIZE 200	      
void CPC_printf(LPCPC_CONFIG pConfig, const char * fmt, ...)
{
    char * buffer=NULL;
    va_list ap;
    va_start(ap,fmt);
    buffer=(char *)CPCAlloc(pConfig,PRINTF_SIZE,MP_WORK);
    vsnprintf(buffer,PRINTF_SIZE-1,fmt,ap);
    va_end(ap);
    pConfig->logConfig.cprint_str(pConfig->logConfig.lpArg,buffer);
    CPCFree(pConfig, buffer, MP_WORK);
}

static char hexchar(int x)
{
    if (x<0 || x>15)
	return '*';
    if (x<10)
	return '0'+x;
    return x-10+'A';	
}

void CPC_hexdump(LPCPC_CONFIG pConfig, const char * msg, void * src, size_t size)
{
    char * buffer=NULL,*ptr=NULL;
    size_t len=strlen(msg);
    size_t sz=len+3*size+size%26+2;
    size_t i;
    unsigned char * s=(unsigned char *)src;
    buffer=(char *)CPCAlloc(pConfig,sz,MP_WORK);
    memcpy(buffer,msg,len);
    ptr=buffer+len;
    for (i=0;i<size;i++)
    {
	*ptr++=hexchar(s[i]>>4);
	*ptr++=hexchar(s[i]&0xF);
	*ptr++=(i%26==25)?'\n':' ';
    }
    *ptr=0;
    pConfig->logConfig.cprint_str(pConfig->logConfig.lpArg,buffer);
    CPCFree(pConfig, buffer, MP_WORK);
}

struct _DRTKTLS_DRV_
{
    CPC_CONFIG Config;
    HCRYPTMODULE hCSP;
    ktls_gost_handle hTLS;
    HCRYPTPROV hProv;
    PRIVKEY Priv; 
    unsigned tlsdb_len;
    unsigned char * our_tlsdb;
};

struct _DRTKTLS_CTX_ 
{    
    PDRTKTLS_DRV pDrv;
    PUBKEY_2012 Pub; 
    CPKTLS_HANDLE hCtx; 
};

static VTABLEPROVSTRUC VTABLE_2001 = { 3, 0, 0, PROV_GOST_2012_256, 0, 0, 0 };

SECURITY_STATUS CAPI_EXTC kinit_gost(PDRTKTLS_DRV pDrv, unsigned * pSize)
{
    ktls_gost_in gi;
    ktls_gost_out go;

    if (pDrv == NULL)
    {
	*pSize=sizeof(DRTKTLS_DRV);
	return STATUS_SUCCESS;
    }
    if (*pSize < sizeof(DRTKTLS_DRV))
    {
	*pSize=sizeof(DRTKTLS_DRV);
	return STATUS_NO_MEMORY;
    }
    memset(pDrv,0,sizeof(DRTKTLS_DRV));

    pDrv->Config.cbSize = sizeof(CPC_CONFIG);
    if (CPCGetDefaultConfig(&pDrv->Config, NULL))
	return STATUS_INTERNAL_ERROR;

    stdInitMemory(&pDrv->Config.pArena, NULL, 0);
    CPCSetLocks(&pDrv->Config);
    pDrv->Config.logConfig.level = 1;
    if (CPCCreateProvider(&pDrv->hCSP, &pDrv->Config))
	return STATUS_INTERNAL_ERROR;

    memset(&gi, 0, sizeof(gi));
    gi.pConfig = &pDrv->Config;
    gi.hCSP = pDrv->hCSP;
    cpktls_init_gost(NULL, &gi, 0, &go);
    gi.allocatedMem = go.requiredMem;
    pDrv->hTLS = (ktls_gost_handle)CPCAlloc(&pDrv->Config, gi.allocatedMem, MP_WORK);
    cpktls_init_gost(pDrv->hTLS, &gi, 0, &go);
    if (pDrv->hCSP->AcquireContext(pDrv->hCSP, &pDrv->hProv, NULL, CRYPT_VERIFYCONTEXT, &VTABLE_2001))
	return STATUS_INTERNAL_ERROR;
    pDrv->our_tlsdb = NULL;
    return STATUS_SUCCESS;
}

SECURITY_STATUS CAPI_EXTC kset_rng(PDRTKTLS_DRV pDrv, unsigned char * pbData, unsigned cbData)
{
    DWORD ret;
    CRYPT_DATA_BLOB cdb;
    cdb.pbData=pbData;
    cdb.cbData=cbData;
    ret = pDrv->hCSP->SetProvParam(pDrv->hCSP,pDrv->hProv,PP_RANDOM,(BYTE *)&cdb,0);
    return (ret==0)?STATUS_SUCCESS:STATUS_INVALID_PARAMETER_1;
}

SECURITY_STATUS CAPI_EXTC kcreate_SA(PDRTKTLS_DRV pDrv)
{
    CPC_printf(&pDrv->Config,"kcreate_SA called\n");
    if (pDrv->our_tlsdb)
    {
	pDrv->hTLS->tlsdb.DestroyPrivKeyFn(0,&pDrv->Priv);
	CPCFree(&pDrv->Config,pDrv->our_tlsdb,MP_WORK);
    }
    if (pDrv->hTLS->tlsdb.CreateEphemFn(pDrv->hCSP,pDrv->hProv,0,&pDrv->Priv,0,&pDrv->tlsdb_len))
	return STATUS_INTERNAL_ERROR;
    CPC_printf(&pDrv->Config,"kcreate_SA ephemfn size determined:%d\n",pDrv->tlsdb_len);
    pDrv->our_tlsdb=(unsigned char *)CPCAlloc(&pDrv->Config,pDrv->tlsdb_len,MP_WORK);
    CPC_printf(&pDrv->Config,"kcreate_SA ptr alloced:%p\n",pDrv->our_tlsdb);
    if (pDrv->hTLS->tlsdb.CreateEphemFn(pDrv->hCSP,pDrv->hProv,0,&pDrv->Priv,pDrv->our_tlsdb,&pDrv->tlsdb_len))
	return STATUS_INTERNAL_ERROR;
    CPC_printf(&pDrv->Config,"kcreate_SA CreateEphomFn ok:\n");
    return STATUS_SUCCESS;
}


SECURITY_STATUS CAPI_EXTC kinit_ctx(PDRTKTLS_DRV pDrv,PDRTKTLS_CTX pDCtx, unsigned * pSize)
{
    if (pDCtx == NULL)
    {
	*pSize=sizeof(DRTKTLS_CTX);
	return STATUS_SUCCESS;
    }
    if (*pSize < sizeof(DRTKTLS_CTX))
    {
	*pSize=sizeof(DRTKTLS_CTX);
	return STATUS_NO_MEMORY;
    }
    memset(pDCtx,0,sizeof(DRTKTLS_CTX));
    pDCtx->pDrv=pDrv;
    return STATUS_SUCCESS;
}

SECURITY_STATUS CAPI_EXTC kdone_ctx(PDRTKTLS_CTX pDCtx)
{
    pDCtx->pDrv->hTLS->tlsdb.DestroyPubKeyFn(0,&pDCtx->Pub);
    if (pDCtx->hCtx && pDCtx->pDrv->hTLS->DeleteContext(pDCtx->pDrv->hTLS,0,pDCtx->hCtx))
	return STATUS_INTERNAL_ERROR;
    check_alloced();	
    return STATUS_SUCCESS;

    // CPCDestroyProvider ??
    return STATUS_SUCCESS;
}

SECURITY_STATUS CAPI_EXTC kdone_gost(PDRTKTLS_DRV pDrv)
{
    pDrv->hTLS->tlsdb.DestroyPrivKeyFn(0,&pDrv->Priv);
    cpktls_shutdown_gost(pDrv->hTLS, 0);
    pDrv->hCSP->ReleaseContext(pDrv->hCSP, pDrv->hProv, 0);
    if (pDrv->our_tlsdb)
	CPCFree(&pDrv->Config, pDrv->our_tlsdb, MP_WORK);
    pDrv->our_tlsdb = 0;
    CPCFree(&pDrv->Config, pDrv->hTLS, MP_WORK);
    pDrv->hTLS = 0;
    if(pDrv->hCSP && pDrv->hCSP->DestroyProvider)
    {
    	pDrv->hCSP->DestroyProvider(pDrv->hCSP);
	pDrv->hCSP = 0;
    }
    if(pDrv->Config.pArena->pDoneMemory)
    {
	CPC_printf(&pDrv->Config,"unfreed:%d\n",totalAlloced);
	pDrv->Config.pArena->pDoneMemory(pDrv->Config.pArena);
    }

    return STATUS_SUCCESS;
}

SECURITY_STATUS CAPI_EXTC kexport_SA(PDRTKTLS_DRV pDrv,unsigned char * pbBuffer,unsigned * puLen)
{
    if (pbBuffer == 0)
    {
	*puLen=pDrv->tlsdb_len;
	return STATUS_SUCCESS;
    }
    else if (*puLen < pDrv->tlsdb_len)
    {
	*puLen= pDrv->tlsdb_len;
	return STATUS_NO_MEMORY;
    }
    memcpy(pbBuffer,pDrv->our_tlsdb,pDrv->tlsdb_len);
    *puLen=pDrv->tlsdb_len;
    return STATUS_SUCCESS;
}

SECURITY_STATUS CAPI_EXTC kimport_SA(PDRTKTLS_CTX pDCtx, unsigned char * pbSA,unsigned cbSA)
{
    CPC_hexdump(&pDCtx->pDrv->Config, "pubkey(compressed):",pbSA,cbSA);
    if (pDCtx->pDrv->hTLS->tlsdb.deSerializePubKeyFn(pDCtx->pDrv->hCSP,pDCtx->pDrv->hProv,pbSA,cbSA,0,&pDCtx->Pub,&pDCtx->pDrv->Config))
	return STATUS_INTERNAL_ERROR;
    CPC_hexdump(&pDCtx->pDrv->Config, "pubkey:",&pDCtx->Pub.bPubKey,pDCtx->Pub.dwPubKeyLen);
    return STATUS_SUCCESS;
}


SECURITY_STATUS CAPI_EXTC kimport_keys(PDRTKTLS_CTX pDCtx,void * pvExpDCtx, unsigned cbExpDCtx)
{
    SecBuffer SB;
    SB.pvBuffer=pvExpDCtx;
    SB.cbBuffer=cbExpDCtx;
    SB.BufferType=SECBUFFER_DATA;
    mark_alloced();
    if (pDCtx->pDrv->hTLS->ImportContext(pDCtx->pDrv->hTLS,&SB,&pDCtx->pDrv->Priv,&pDCtx->Pub,0,&pDCtx->hCtx))
	return STATUS_INTERNAL_ERROR;
    return STATUS_SUCCESS;	
}

SECURITY_STATUS CAPI_EXTC kexport_keys(PDRTKTLS_CTX pDCtx, void * pvExpDCtx, unsigned *pcbExpDCtx)
{
    SecBuffer SB;
    DWORD dwErr;

    SB.pvBuffer=pvExpDCtx;
    SB.cbBuffer=*pcbExpDCtx;
    SB.BufferType=SECBUFFER_DATA;
    dwErr = pDCtx->pDrv->hTLS->ExportContext(pDCtx->pDrv->hTLS, pDCtx->hCtx,
					&pDCtx->pDrv->Priv, &pDCtx->Pub,
					CPKTLS_CONTEXT_EXPORT_TO_USER|
					CPKTLS_CONTEXT_EXPORT_DELETE_OLD,
					&SB);
    *pcbExpDCtx=SB.cbBuffer;
    return STATUS_SUCCESS;
}


SECURITY_STATUS CAPI_EXTC kreadall(PDRTKTLS_CTX pDCtx, int Socket, void * pvExtra, unsigned *pcbExtra)
{
    //
    // Чтение данных от сервера до их окончания.
    //
    SECURITY_STATUS scRet;
    SecPkgContext_StreamSizes Sizes;
    SecBufferDesc   Message;
    SecBuffer       Buffers[4];
    SecBuffer *     pDataBuffer;
    SecBuffer *     pExtraBuffer;
    BYTE * pbIoBuffer = NULL;
    DWORD cbIoBufferLength = 0;
    DWORD cbIoBuffer;
    DWORD cbData=0;
    int i;

    scRet = pDCtx->pDrv->hTLS->QueryAttributes(pDCtx->pDrv->hTLS,
				   pDCtx->hCtx, 
                                   SECPKG_ATTR_STREAM_SIZES,
                                   &Sizes);
    if(scRet != SEC_E_OK)
    {
        CPC_printf(&pDCtx->pDrv->Config,"**** Error 0x%x reading SECPKG_ATTR_STREAM_SIZES\n", scRet);
        return scRet;
    }

    CPC_printf(&pDCtx->pDrv->Config,"\nHeader: %d, Trailer: %d, MaxMessage: %d\n",
        Sizes.cbHeader,
        Sizes.cbTrailer,
        Sizes.cbMaximumMessage);

    //
    // Выделение рабочего буфера. Открытый текст, передаваемый EncryptMessage,
    // должен ен провосходить по размерам 'Sizes.cbMaximumMessage', поэтому 
    // размер буфера равен этой величене, сложеной с размерами заголовка и 
    // заключительной части.
    // 
    cbIoBufferLength = Sizes.cbHeader + 
                       Sizes.cbMaximumMessage +
                       Sizes.cbTrailer;// Могут быть такие реализации 
    //сервера и клиента что длина буфера будет cbIoBufferLength
    // вычисленная таким образом + 2048 байт см. RFC 2246

    pbIoBuffer = CPCAlloc(&pDCtx->pDrv->Config,cbIoBufferLength,MP_WORK);
    if(pbIoBuffer == NULL)
    {
        CPC_printf(&pDCtx->pDrv->Config,"**** Out of memory (2)\n");
        return SEC_E_INTERNAL_ERROR;
    }


    cbIoBuffer = 0;

    for(;;)
    {
        //
        // Чтение данных.
        //

        if(0 == cbIoBuffer || scRet == SEC_E_INCOMPLETE_MESSAGE)
        {
            cbData = emulate_recv(Socket, 
                          pbIoBuffer + cbIoBuffer, 
                          cbIoBufferLength - cbIoBuffer, 
                          0);
            if(cbData == (DWORD)(-1))
            {
#ifdef KRN_EMUL	    
                CPC_printf(&pDCtx->pDrv->Config,"**** Error %d reading data from server\n", errno);
#endif		
                return SEC_E_INTERNAL_ERROR;
            }
            else if(cbData == 0)
            {
                // Произошел разрыв соединения с сервером.
                if(cbIoBuffer)
                {
                    CPC_printf(&pDCtx->pDrv->Config,"**** Server unexpectedly disconnected\n");
                    scRet = SEC_E_INTERNAL_ERROR;
		    goto ret_scRet;
                }
                else
                {
                    scRet= SEC_E_OK;
		    goto ret_scRet;
                }
            }
            else
            {
                CPC_printf(&pDCtx->pDrv->Config,"%d bytes of (encrypted) application data received\n", cbData);
		CPC_hexdump(&pDCtx->pDrv->Config,"Recieved:\n",pbIoBuffer+cbIoBuffer,cbData); 

                
                cbIoBuffer += cbData;
            }
        }
        // 
        // Попытка расшифрования поченных от сервера данных.
        //

        Buffers[0].pvBuffer     = pbIoBuffer;
        Buffers[0].cbBuffer     = cbIoBuffer;
        Buffers[0].BufferType   = SECBUFFER_DATA;

        Buffers[1].BufferType   = SECBUFFER_EMPTY;
        Buffers[2].BufferType   = SECBUFFER_EMPTY;
        Buffers[3].BufferType   = SECBUFFER_EMPTY;

        Message.ulVersion       = SECBUFFER_VERSION;
        Message.cBuffers        = 4;
        Message.pBuffers        = Buffers;

        scRet = pDCtx->pDrv->hTLS->UnSeal(pDCtx->pDrv->hTLS,pDCtx->hCtx, &Message, 0, 0);
	// SSPDecryptMessage

        if(scRet == SEC_E_INCOMPLETE_MESSAGE)
        {
            // Входной буфер содержит только фрагмент зашифрованных данных.
            // Продолжение цикла и чтение данных.
            continue;
        }

        // Сервер завершил сессию
        if(scRet == SEC_I_CONTEXT_EXPIRED)
            break;

        if( scRet != SEC_E_OK && 
            scRet != SEC_I_RENEGOTIATE && 
            scRet != SEC_I_CONTEXT_EXPIRED)
        {
            CPC_printf(&pDCtx->pDrv->Config,"**** Error 0x%x returned by DecryptMessage\n", scRet);
	    goto ret_scRet;
        }

        // Расположение данных и (необязательно) добавление буферов.
        pDataBuffer  = NULL;
        pExtraBuffer = NULL;
        for(i = 1; i < 4; i++)
        {

            if(pDataBuffer == NULL && Buffers[i].BufferType == SECBUFFER_DATA)
            {
                pDataBuffer = &Buffers[i];
                CPC_printf(&pDCtx->pDrv->Config,"Buffers[%d].BufferType = SECBUFFER_DATA\n",i);
            }
            if(pExtraBuffer == NULL && Buffers[i].BufferType == SECBUFFER_EXTRA)
            {
                pExtraBuffer = &Buffers[i];
            }
        }

        // Вывод или обработка расшифрованных данных.
        if(pDataBuffer)
        {
            CPC_printf(&pDCtx->pDrv->Config,"Decrypted data size: %d bytes\n", pDataBuffer->cbBuffer);
            CPC_printf(&pDCtx->pDrv->Config,"Decrypted data: %.*s\n", pDataBuffer->cbBuffer,pDataBuffer->pvBuffer);
        }

        // Перенос всех "extra" данных во входной буфер.
        if(pExtraBuffer)
        {
            memmove(pbIoBuffer, pExtraBuffer->pvBuffer, pExtraBuffer->cbBuffer);
            cbIoBuffer = pExtraBuffer->cbBuffer;
        }
        else
        {
            cbIoBuffer = 0;
        }
	if (scRet == SEC_I_RENEGOTIATE)
	{
	    if (cbIoBuffer)
	    {
		memcpy(pvExtra,pbIoBuffer,cbIoBuffer);	
		*pcbExtra=cbIoBuffer;
	    }
	    else
	    {
		*pcbExtra=0;
       	    }
	    goto ret_scRet;
	}
 
    }
ret_scRet:
    CPCFree(&pDCtx->pDrv->Config, pbIoBuffer, MP_WORK);
    return scRet;	
}    
