/*
 * Copyright(C) 2000-2010 Проект ИОК
 *
 * Этот файл содержит информацию, являющуюся
 * собственностью компании Крипто Про.
 *
 * Любая часть этого файла не может быть скопирована,
 * исправлена, переведена на другие языки,
 * локализована или модифицирована любым способом,
 * откомпилирована, передана по сети с или на
 * любую компьютерную систему без предварительного
 * заключения соглашения с компанией Крипто Про.
 *
 * Программный код, содержащийся в этом файле, предназначен
 * исключительно для целей обучения и не может быть использован
 * для защиты информации.
 *
 * Компания Крипто-Про не несет никакой
 * ответственности за функционирование этого кода.
 */


#include "stdafx.h"
#include "UnixEnroll.h"
#include "UnixRequest.h"
#include "CPCA15UserInfo.h"
#include "CPCA20UserInfo.h"
#include "CPCA20Request.h"
#include <vector>
#include <wchar.h>
#include <time.h> 
#include <cassert>
#include <cstdio>
#include "BSTR.h"
#include "compiler_attributes.h"

#ifdef _WIN32
#   define CONTAINER L"\\\\.\\REGISTRY\\CPCA20test"
#else
#   define CONTAINER L"\\\\.\\HDIMAGE\\CPCA20test"
#endif
#define CERFILE "CPCA20test.cer"
#define MAX_SUBMISSIONS_COUNT 10
#define REQUEST_ID_LEN        37
#define INTERNAL_ERROR_CODE   15

enum CertificateSubmissionStatus {
    SUCCESS_CERTIFICATE_GET = 0,
    PENDING_CERTIFICATE,
    ERROR_CERTIFICATE_GET
};

//--------------------------------------------------------------------
// Пример кода взаимодействия - регистрации пользователя, создания 
// запроса и его отправки - с УЦ 20

class Callbacks: public UnixEnroll::UserCallbacks
{
public:
    bool askPermissionToAddToRootStore( const BYTE* pbCert, DWORD cbCert, bool force = false) const;
    UserCallbacks* clone() const;
    virtual ~Callbacks() {}
};

//--------------------------------------------------------------------
// Функция askPermissionToAddToRootStore всегда будет разрешать 
// записывать сертификат в корневое хранилище

bool Callbacks::askPermissionToAddToRootStore( const BYTE* pbCert, DWORD cbCert, bool) const
{
    int res;
    DWORD size = 0;
    PBYTE pHashValue = NULL;
    res = CryptHashCertificate(0,0,0,pbCert,cbCert,NULL,&size);
    if(!res) {
	return false;
    }

    pHashValue =(BYTE *)malloc(size);
    if(!pHashValue) {
	return false;
    }

    res = CryptHashCertificate(0,0,0,pbCert,cbCert,pHashValue,&size);
    if(!res) {
	free(pHashValue);
	return false;
    }
    free(pHashValue);
    return true;
}

UnixEnroll::UserCallbacks* Callbacks::clone() const
{
    return new Callbacks();
}

static void ATTR_NORETURN HandleError(HRESULT err, const char *s)
{
    ::printf("Error number     : 0x%x\n", err);
    ::printf("Error description: %s\n", s);
    if(!err)
	err = 1;
    ::exit(err);
}

static void ATTR_NORETURN HandleLastError(const char *s)
{
    const HRESULT err = (HRESULT) GetLastError();
    HandleError(err, s);
}

static HRESULT addTemlateExtension( std::string  const& ext, CPEnroll *pEnroll4 )
{
    HRESULT hr = -1;
    std::vector<char> encodedExt( ext.begin(), ext.end() );
    encodedExt.push_back( 0 );

    BSTR value = 0;
    CERT_NAME_VALUE templateName;
    templateName.dwValueType = CERT_RDN_UTF8_STRING;
    templateName.dwValueType = CERT_RDN_IA5_STRING;
    templateName.Value.cbData = encodedExt.size();
    templateName.Value.pbData = ( BYTE* )&encodedExt[0];

    DWORD size = 0;
    if ( !CryptEncodeObject( X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
	X509_ANY_STRING, &templateName, 0, &size ) )
    {
	return GetLastError();
    }
    std::vector<BYTE> encodedNameSeq( size );
    if ( !CryptEncodeObject( X509_ASN_ENCODING | PKCS_7_ASN_ENCODING,
	X509_ANY_STRING, &templateName, &encodedNameSeq[0], &size ) )
    {
	return GetLastError();
    }

    CERT_EXTENSION tmplateExtension;
    tmplateExtension.fCritical = 0;
    std::string tmp( szOID_ENROLL_CERTTYPE_EXTENSION );
    tmplateExtension.pszObjId = &tmp[0];

    hr = ConvertBinToBSTR(encodedNameSeq, &value);
    if(hr != S_OK)
        return hr;

    BSTR oid = ConvertStringToBSTR( tmplateExtension.pszObjId );
    hr  = pEnroll4->addExtensionToRequest(tmplateExtension.fCritical, oid, value);

    SysFreeString( oid );
    SysFreeString( value );
    return hr;
}

static long create_request(BSTR bstrRDN, BSTR bstrEKUsage, BSTR contName, char** chTmp)
{
    HRESULT res = -1;
    Callbacks tb;
    BSTR bstrTemplate = SysAllocString( L"User" );
    DWORD dwKeySpec = AT_KEYEXCHANGE;
    BSTR strRequest = 0;
    UnixEnroll * pEnroll = new UnixEnroll(tb, false);
    /*hr = addTemlateExtension("1.2.643.2.2.46.0.7", pEnroll);*/

    // put_ContainerName создает имя контейнера
    pEnroll->put_ContainerName(contName);
    //задаём дефолтный пин "1" на контейнер 
    std::string PIN = "1";
    CSecurePin pin(PIN.length() + 1);
    memcpy(pin.ptr_rw(), PIN.c_str(), PIN.length() + 1);
    HRESULT hr = pEnroll->put_PIN(pin);
    if(hr)
        HandleError(hr, "Error in put_Pin.");
    pin.clean();
    // put_LimitExchangeKeyToEncipherment задает ограничения ключа AT_KEYEXCHANGE, false
    // и шифрование, и подпись, true - только шифрование
    pEnroll->put_LimitExchangeKeyToEncipherment(TRUE);
    // put_KeySpec задает тип закрытого ключа
    pEnroll->put_KeySpec(dwKeySpec);
    // put_ProviderType задает тип CSP для создания ключа.
    pEnroll->put_ProviderType(PROV_GOST_2001_DH);
    // put_RequestStoreFlags задаеь хранилище запроса и флаги его открытия.
    pEnroll->put_RequestStoreFlags(CERT_SYSTEM_STORE_CURRENT_USER);

    char* tp = ConvertBSTRToString( bstrTemplate );
    // происходит создание запроса функцией createRequest
    res = addTemlateExtension( tp, pEnroll );
    if ( res != S_OK )
    {
	SysFreeString(bstrTemplate);
	delete pEnroll;
	return res;
    }

    res = pEnroll->createRequest(XECR_PKCS10_V2_0, bstrRDN, bstrEKUsage, &strRequest);
    if(res == S_OK) {
	*chTmp = ConvertBSTRToString(strRequest);
	SysFreeString(strRequest);
	printf("charRequest = %s\n", *chTmp);
    }

    delete [] tp;
    delete pEnroll;
    SysFreeString(bstrTemplate);
    return res;
}

static CertificateSubmissionStatus ProcessCertificateDisposition(LONG pDisposition)
{
    switch (pDisposition) {
        case CR_DISP_DENIED:
	    printf("request was denied. \n");
	    break;

	case CR_DISP_ERROR:
	    printf("error request. \n");
	    break;

	case CR_DISP_INCOMPLETE:
	    printf("request was incompleted\n");
	    break;

	case CR_DISP_ISSUED:
	    printf("certificate was issued. \n");
            return SUCCESS_CERTIFICATE_GET;
	    break;

	case CR_DISP_ISSUED_OUT_OF_BAND:
	    //Сертификат выдан (будет выдан) непосредственно на УЦ, а не по сети
	    printf("certificate was issued out of band. \n");
	    break;

	case CR_DISP_UNDER_SUBMISSION:
            return PENDING_CERTIFICATE;
            break;	
	default:
	    printf("default. \n");
	    break;
    }
    return ERROR_CERTIFICATE_GET;
}

static HRESULT LoadAndSaveCert(UnixRequest *pRequest)
{
    BSTR pMyCert1;
    const HRESULT res=pRequest->GetCertificate(CR_OUT_BASE64HEADER, &pMyCert1);
    if (res!= S_OK)
    {
        SysFreeString(pMyCert1);
        return res;
    }

    char *pCertStr1 = ConvertBSTRToString(pMyCert1);
    SysFreeString(pMyCert1);

    //пишем сертификат в дефолтный файл
    FILE *fp = fopen(CERFILE, "wb");
    if (fp == NULL)
        HandleLastError("Error while openning file.");
    fputs(pCertStr1, fp);
    fclose(fp);

    //и заодно выводим его на экран
    printf("Cert:\n%s", pCertStr1);
    delete [] pCertStr1;
    return S_OK;
}

int main(int argc, char* argv[])
{    
    std::string strDN="2.5.4.3=";
    CPCA20UserFields ui;
    HRESULT res(0);
    LONG Flags = CR_IN_BASE64 | CR_IN_PKCS10;
    LONG pDisposition;
    BSTR bstrRDN = 0;
    BSTR bstrEKUsage = SysAllocString( L"1.3.6.1.5.5.7.3.4,1.2.643.2.2.34.6,1.3.6.1.5.5.7.3.2" );
    BSTR bstrRequest = 0;
    BSTR bstrTemplate = SysAllocString( L"User" );
    BSTR bstrTokenID = 0;
    BSTR UrlOfCA;
    std::string UIURL;
    CertificateSubmissionStatus submissionStatus;
    int counter = 0;
    bool isUnderSubmission = false; 
    uintptr_t rID = 0;
    BSTR bstrRequestId;
    BSTR bsrtUIURL;
    CPCA20Request* CPCA20R;
    LONG userRegisterStatus;

    if (argc == 5 && !strcmp(argv[1], "-serv") && !strcmp(argv[3], "-folder"))
    {
        char *tmpChar = argv[2];
        if (tmpChar[strlen(tmpChar) - 1] != '/') {
            strcat(tmpChar, "/");
        }

	BSTR prefix = ConvertStringToBSTR(tmpChar);
	BSTR suffix = ConvertStringToBSTR(argv[4]);
	UrlOfCA = SysAllocStringLen(L"", SysStringLen(prefix) + SysStringLen(suffix));
	wcscpy(UrlOfCA, prefix);
	wcscat(UrlOfCA, suffix);
	SysFreeString(prefix);
	SysFreeString(suffix);

        UIURL = tmpChar;
    } else {
        UrlOfCA =  SysAllocString(L"https://ca20-x64-w12r2.cp.ru/ui/MainRA");
        UIURL = "https://ca20-x64-w12r2.cp.ru/ui/";
    }
    bsrtUIURL = ConvertStringToBSTR(UIURL.c_str());
    // адрес УЦ с папкой нужен для вызовов GetUserRegisterInfo и RegisterUser
    BSTR contName = SysAllocString(CONTAINER);
    std::string CAId = "CPCA20";
    // адрес УЦ без папки
    char* charRequest = NULL;
    UnixRequest * pRequest = UnixRequest::URFactory( CAId.c_str() );
    if(!pRequest) 
	HandleLastError("Error creating prequest object!\n");

    //-----------------------------------------------------------------------------------------//
    // устнавовка методов аутентификации при взаимодействии с УЦ и получение данных формы регистрации
    res = pRequest->SetCredential( NULL, X509AuthAnonymous, X509CC_TLS, NULL, NULL );
    if ( res != S_OK )
	HandleError( res, "Error Set Credential!\n" );

    res = pRequest->GetUserRegisterInfo( UrlOfCA, &ui );
    if ( res != S_OK )
	HandleError( res, "Error GetUserRegisterInfo!\n" );
    //-----------------------------------------------------------------------------------------//

    // заполнение обязательных полей случайными значениями
    {
	HCRYPTPROV hProv = 0;
	if (!CryptAcquireContext(&hProv, NULL, NULL, PROV_GOST_2012_256, CRYPT_VERIFYCONTEXT)) {
	    HandleLastError("Error during CryptAcquireContext.");
	}
	unsigned seed;
	BOOL bRet = CryptGenRandom(hProv, sizeof(seed), (BYTE *)&seed);
	if (hProv) {
	    CryptReleaseContext(hProv, 0);
	}
	if (!bRet) {
	    HandleLastError("Random bytes were not correctly generated.");
	}
	srand(seed);
    }

    int flag = 1;
    for ( CPCA20UserData::iterator it = ui.fields.begin();it != ui.fields.end();it++ )
    {
	if (it->oid.compare("1.3.6.1.4.1.311.20.2.3") == 0) {
	  //Для UPN
	  it->value = "test@cryptopro.ru"; 
	}
	if (it->oid.compare("2.5.4.3") == 0) {
	    for (int i = 1; i <= 9; i++)
	    {
		flag = rand() % 2;
		if (flag)
		    it->value += ('a' + rand() % 26);
		else
		    it->value += ('0' + rand() % 10);
	    }
	    flag = 0;
	}
    }

    if (flag) {
	printf("Oid 2.5.4.3 not found.\n");
    }

    // регистрация пользователя
    res = pRequest->RegisterUser( UrlOfCA, &ui );
    if ( res != S_OK )
	HandleError( res, "Error Register User!\n" );

    switch ( ui.status )
    {
    case UR_DISP_ISSUED:
	printf("ISSUED\n");
	break;
    case UR_DISP_UNDER_SUBMISSION:
	printf("UNDER_SUBMISSION\n");
	break;
    default:
	printf("CANT_GET_REG_STATUS\n");
	break;
    }

    printf("TokenID = %s\n\n", ui.TokenID.c_str());
    strDN+=ui.TokenID;
    printf("RDN = %s\n\n", strDN.c_str());
    bstrTokenID = ConvertStringToBSTR(ui.TokenID.c_str());
    printf("sbPassword = %s\n\n", (*(ui.sbPassword)).ptr_rw());

    res = pRequest->SetCredential( NULL,
	X509AuthUsername,
	X509CC_TLS,
	bstrTokenID,
	ui.sbPassword );
    if ( res != S_OK )
	HandleError( res, "Error Set Credential!\n" );

    counter = 0;
    userRegisterStatus = ui.status;
    while (userRegisterStatus == UR_DISP_UNDER_SUBMISSION && counter < MAX_SUBMISSIONS_COUNT) {
        res = pRequest->GetUserRegisterStatus(bsrtUIURL, 0, &userRegisterStatus);
        if (res != S_OK)
            HandleError( res, "Error checking user register status.");
        ++counter;
        if (UR_DISP_UNDER_SUBMISSION == userRegisterStatus)
            Sleep(1000);
    }
    if (userRegisterStatus == UR_DISP_ISSUED)
	printf("User registered on iteration %d.\n", counter);
    if (userRegisterStatus == UR_DISP_UNDER_SUBMISSION)
	printf("User registration is still under submission.\n");

    SysFreeString(UrlOfCA); // заполняли раньше
    UrlOfCA = ConvertStringToBSTR(UIURL.c_str());
    bstrRDN = ConvertStringToBSTR(strDN.c_str());
    res = create_request(bstrRDN, bstrEKUsage, contName, &charRequest);

    if (res != S_OK)
	HandleError( res, "Error create request.");
    printf("Successful create request.\n");

    //-----------------------------------------------------------------------------------------//
    // отправка и обработка результата запроса
    bstrRequest = ConvertStringToBSTR(charRequest);
    printf("charRequest = %s\n", charRequest);

    res = pRequest->Submit(Flags, bstrRequest, bstrTemplate, UrlOfCA, &pDisposition);

    if (res != S_OK)
	HandleError( res, "Error submit." );

    submissionStatus = ProcessCertificateDisposition(pDisposition);
    switch(submissionStatus) {
	case SUCCESS_CERTIFICATE_GET:
	    res = LoadAndSaveCert(pRequest);
	    if (res != S_OK) {
		res = INTERNAL_ERROR_CODE;
		goto er;
	    }
	    break;
	case PENDING_CERTIFICATE:
	    isUnderSubmission = true;
	    break;
	case ERROR_CERTIFICATE_GET:
	    break;
	default:
	    break;
    }
    if (!isUnderSubmission)
	printf("Successful submitting. \n");

    CPCA20R = static_cast<CPCA20Request*>(pRequest);
    res = CPCA20R->GetRequestStrId(&rID);
    if (res != S_OK)
	HandleError(res, "Error getting request id");
    char requestId[REQUEST_ID_LEN];
    memcpy(requestId, (void *)rID, sizeof(requestId));

    bstrRequestId = ConvertStringToBSTR(requestId);
    BSTR fullPendingURL;
    fullPendingURL = SysAllocStringLen(L"", SysStringLen(bsrtUIURL) + SysStringLen(bstrRequestId));
    wcscpy(fullPendingURL, bsrtUIURL);
    wcscat(fullPendingURL, bstrRequestId);
    SysFreeString(bsrtUIURL);
    SysFreeString(bstrRequestId);

    if (isUnderSubmission) {
        counter = 0;
        while (isUnderSubmission && counter < MAX_SUBMISSIONS_COUNT) {
            Sleep(1000);
            ++counter;
            isUnderSubmission = false;
            res = pRequest->RetrievePending(0 /*pRequestId*/, fullPendingURL, &pDisposition);
            if (res != S_OK)
        	HandleError( res, "Error getting certificate status." );
            submissionStatus = ProcessCertificateDisposition(pDisposition);
            switch(submissionStatus) {
                case SUCCESS_CERTIFICATE_GET:
                    res = LoadAndSaveCert(pRequest);
                    if (res != S_OK) {
                        res = INTERNAL_ERROR_CODE;
                    }
                    break;
                case PENDING_CERTIFICATE:
                    isUnderSubmission = true;
                    break;
                case ERROR_CERTIFICATE_GET:
                    break;
                default:
                    break;
            }
        }
        if (!isUnderSubmission)
            printf("Successful submitting on iteration %d.\n", counter);
        else 
            printf("Certificate is under submission.\n");
    }
    printf("RequestId = %s\n", requestId);
    SysFreeString(fullPendingURL);

er:
    delete pRequest;
    SysFreeString( UrlOfCA );
    SysFreeString( contName );
    SysFreeString( bstrTemplate );
    SysFreeString( bstrRDN );
    SysFreeString( bstrEKUsage );
    SysFreeString( bstrRequest );
    SysFreeString( bstrTokenID );
    delete [] charRequest;
    return res;
}
