#include <windows.h>

#include "wsocket.h"
#include "sslsocket.h"

//  .
#define SOCKET_IO_BUFFER 0x10000

//  InitializeSecurityContext (    IE).
#define ISC_FLAGS (ISC_REQ_ALLOCATE_MEMORY | ISC_REQ_CONFIDENTIALITY | ISC_REQ_EXTENDED_ERROR | ISC_REQ_MANUAL_CRED_VALIDATION |  ISC_REQ_REPLAY_DETECT | ISC_REQ_SEQUENCE_DETECT | ISC_REQ_STREAM)

static void safeFreeContextBuffer(SecBuffer *buffer)
{
  if(buffer->pvBuffer != NULL && buffer->cbBuffer != 0)CWA(secur32, FreeContextBuffer)(buffer->pvBuffer);
  buffer->pvBuffer   = NULL;
  buffer->cbBuffer   = 0;
  buffer->BufferType = SECBUFFER_EMPTY;
}

void SslSocket::init(void)
{

}

void SslSocket::uninit(void)
{

}

SslSocket::SERVERDATA *SslSocket::_createServerData(PCCERT_CONTEXT context, DWORD enabledProtocols)
{
  SERVERDATA sd;
  SCHANNEL_CRED sc;

  sc.dwVersion               = SCHANNEL_CRED_VERSION;
  sc.cCreds                  = 1;
  sc.paCred                  = &context;
  sc.hRootStore              = NULL;
  sc.cMappers                = 0;
  sc.aphMappers              = NULL;
  sc.cSupportedAlgs          = 0;
  sc.palgSupportedAlgs       = NULL;
  sc.grbitEnabledProtocols   = enabledProtocols;
  sc.dwMinimumCipherStrength = 0;
  sc.dwMaximumCipherStrength = 0;
  sc.dwSessionLifespan       = (MAXDWORD - 1);
  sc.dwFlags                 = SCH_CRED_NO_SYSTEM_MAPPER;
  sc.dwCredFormat            = 0;

  if(CWA(secur32, AcquireCredentialsHandleW)(NULL, UNISP_NAME_W, SECPKG_CRED_INBOUND, NULL, &sc, NULL, NULL, &sd.credHandle, NULL) == SEC_E_OK)
  {
    sd.flags = 0;
    
    SERVERDATA *newSd = (SERVERDATA *)Mem::copyEx(&sd, sizeof(SERVERDATA));
    if(newSd != NULL)return newSd;
    CWA(secur32, FreeCredentialsHandle)(&sd.credHandle);
  }
  return NULL;
}

void SslSocket::_freeServerData(SERVERDATA *sd)
{
  CWA(secur32, FreeCredentialsHandle)(&sd->credHandle);
  Mem::free(sd);
}

bool SslSocket::_serverHandshake(SERVERDATA *sd, void *recv, DWORD recvSize, void *send, DWORD *sendSize)
{
  SecBufferDesc inBuffer;
  SecBufferDesc outBuffer;
  SecBuffer     inBuffers[2];
  SecBuffer     outBuffers[2];
  DWORD         contextAttr;

  //recv  .
  inBuffers[0].pvBuffer    = (void *)recv;
  inBuffers[0].cbBuffer    = recvSize;
  inBuffers[0].BufferType  = SECBUFFER_TOKEN;

  inBuffers[1].pvBuffer    = NULL;
  inBuffers[1].cbBuffer    = 0;
  inBuffers[1].BufferType  = SECBUFFER_EMPTY;

  inBuffer.cBuffers        = 2;
  inBuffer.pBuffers        = inBuffers;
  inBuffer.ulVersion       = SECBUFFER_VERSION;
  
  //send .
  outBuffers[0].pvBuffer   = NULL;
  outBuffers[0].BufferType = SECBUFFER_TOKEN;
  outBuffers[0].cbBuffer   = 0;

  outBuffer.cBuffers       = 1;
  outBuffer.pBuffers       = outBuffers;
  outBuffer.ulVersion      = SECBUFFER_VERSION;

  // .
/*  SECURITY_STATUS status = CWA(secur32, AcceptSecurityContext)(&sd->credHandle,
                                                               sd->flags & SDF_STATUS_HANDSHAKE ? &sd->ctxtHandle : NULL,
                                                               &inBuffer,
                                                               ASC_REQ_ALLOCATE_MEMORY | ASC_REQ_EXTENDED_ERROR | ASC_REQ_REPLAY_DETECT | ASC_REQ_SEQUENCE_DETECT | ASC_REQ_STREAM,
                                                               &outBuffer,
                                                               &sd->ctxtHandle, //On subsequent calls, phNewContext can be the same as the handle specified in the phContext parameter.
                                                               &contextAttr,
                                                               NULL
                                                              );
  */
  sd->flags |= SDF_STATUS_HANDSHAKE;



}

/*bool SslSocket::_startServerEncryption(SOCKET socket, SERVERDATA *sd, DWORD enabledProtocols, DWORD timeout)
{
  AcceptSecurityContext(&sd->credHandle, NULL, );


}
*/
bool SslSocket::_startClientEncryption(SOCKET socket, LPWSTR serverName, SOCKETDATA *sd, DWORD enabledProtocols, DWORD timeout)
{
  bool ok = false;
  Mem::_zero(sd, sizeof(SOCKETDATA));

  //.
  {
    SCHANNEL_CRED sc;

    sc.dwVersion               = SCHANNEL_CRED_VERSION;
    sc.cCreds                  = 0;
    sc.paCred                  = NULL;
    sc.hRootStore              = NULL;
    sc.cMappers                = 0;
    sc.aphMappers              = NULL;
    sc.cSupportedAlgs          = 0;
    sc.palgSupportedAlgs       = NULL;
    sc.grbitEnabledProtocols   = enabledProtocols;
    sc.dwMinimumCipherStrength = 0;
    sc.dwMaximumCipherStrength = 0;
    sc.dwSessionLifespan       = (MAXDWORD - 1);
    sc.dwFlags                 = SCH_CRED_MANUAL_CRED_VALIDATION | SCH_CRED_NO_DEFAULT_CREDS | SCH_CRED_NO_SERVERNAME_CHECK;
    sc.dwCredFormat            = 0;

    if(CWA(secur32, AcquireCredentialsHandleW)(NULL, UNISP_NAME_W, SECPKG_CRED_OUTBOUND, NULL, &sc, NULL, NULL, &sd->sh, NULL) != SEC_E_OK)return false;
  }

  DWORD dwContextAttr;
  SecBufferDesc OutBuffer;
  SecBuffer OutBuffers[1];

  OutBuffers[0].pvBuffer   = NULL;
  OutBuffers[0].BufferType = SECBUFFER_TOKEN;
  OutBuffers[0].cbBuffer   = 0;

  OutBuffer.cBuffers  = 1;
  OutBuffer.pBuffers  = OutBuffers;
  OutBuffer.ulVersion = SECBUFFER_VERSION;

  // handshake.
  SECURITY_STATUS ssFirst = CWA(secur32, InitializeSecurityContextW)(&sd->sh, NULL, serverName, ISC_FLAGS, 0, 0, NULL, 0, &sd->ch, &OutBuffer, &dwContextAttr, NULL);

  #if defined(WDEBUG1)
  WDEBUG1(WDDT_INFO, "ssFirst=0x%08X", ssFirst);
  #endif

  if(ssFirst == SEC_I_CONTINUE_NEEDED && WSocket::tcpSend(socket, OutBuffers[0].pvBuffer, OutBuffers[0].cbBuffer))
  {
    safeFreeContextBuffer(&OutBuffers[0]);

    //    .
    LPBYTE recvBuffer = (LPBYTE)Mem::alloc(SOCKET_IO_BUFFER);
    DWORD recvSize   = 0;

    if(recvBuffer != NULL)
    {
      SecBufferDesc   InBuffer;
      SecBuffer       InBuffers[2];
      SECURITY_STATUS ssLoop = SEC_I_CONTINUE_NEEDED;

      while(ssLoop == SEC_I_CONTINUE_NEEDED || ssLoop == SEC_E_INCOMPLETE_MESSAGE)
      {
        //   .
        if(recvSize == 0 || ssLoop == SEC_E_INCOMPLETE_MESSAGE)
        {
          int r = WSocket::tcpRecv(socket, recvBuffer + recvSize, SOCKET_IO_BUFFER - recvSize, timeout);
          if(r == 0 || r == SOCKET_ERROR)break;
          recvSize += r;
        }

        // .
        InBuffers[0].pvBuffer    = (void *)recvBuffer;
        InBuffers[0].cbBuffer    = recvSize;
        InBuffers[0].BufferType  = SECBUFFER_TOKEN;

        InBuffers[1].pvBuffer    = NULL;
        InBuffers[1].cbBuffer    = 0;
        InBuffers[1].BufferType  = SECBUFFER_EMPTY;

        InBuffer.cBuffers        = 2;
        InBuffer.pBuffers        = InBuffers;
        InBuffer.ulVersion       = SECBUFFER_VERSION;

        OutBuffers[0].pvBuffer   = NULL;
        OutBuffers[0].BufferType = SECBUFFER_TOKEN;
        OutBuffers[0].cbBuffer   = 0;

        OutBuffer.cBuffers       = 1;
        OutBuffer.pBuffers       = OutBuffers;
        OutBuffer.ulVersion      = SECBUFFER_VERSION;

        dwContextAttr = 0;
        ssLoop = CWA(secur32, InitializeSecurityContextW)(&sd->sh, &sd->ch, serverName, ISC_FLAGS, 0, 0, &InBuffer, 0, NULL, &OutBuffer, &dwContextAttr, NULL);

        #if defined(WDEBUG1)
        WDEBUG1(WDDT_INFO, "ssLoop=0x%08X", ssLoop);
        #endif

        if(ssLoop == SEC_E_INCOMPLETE_MESSAGE)continue;
        if(ssLoop == SEC_E_OK || ssLoop == SEC_I_CONTINUE_NEEDED || ((FAILED(ssLoop)) && (dwContextAttr & ISC_RET_EXTENDED_ERROR)))
        {
          // .
          if(OutBuffers[0].cbBuffer != 0 && OutBuffers[0].pvBuffer != NULL)
          {
            if(!WSocket::tcpSend(socket, OutBuffers[0].pvBuffer, OutBuffers[0].cbBuffer))break;
            safeFreeContextBuffer(&OutBuffers[0]);
          }

          //    .
          if(InBuffers[1].BufferType == SECBUFFER_EXTRA)
          {
            #if defined(WDEBUG1)
            WDEBUG1(WDDT_INFO, "SECBUFFER_EXTRA detected with size %u.", InBuffers[1].cbBuffer);
            #endif

            LPBYTE pOffset = recvBuffer + recvSize - InBuffers[1].cbBuffer;

            if(ssLoop == SEC_E_OK)
            {
              sd->extraBuffer = Mem::copyEx(pOffset, InBuffers[1].cbBuffer);
              if(sd->extraBuffer == NULL)break;
              sd->extraBufferSize = InBuffers[1].cbBuffer;
            }
            else
            {
              Mem::_copy(recvBuffer, pOffset, InBuffers[1].cbBuffer);
              recvSize = InBuffers[1].cbBuffer;
            }
          }
          else recvSize = 0;

          if(ssLoop == SEC_E_OK)
          {
            #if defined(WDEBUG1)
            WDEBUG1(WDDT_INFO, "dwContextAttr=0x%08X", dwContextAttr);
            #endif

            ok = true;
          }
        }
      }
      Mem::free(recvBuffer);
    }
  }

  if(ok && CWA(secur32, QueryContextAttributesW)(&sd->ch, SECPKG_ATTR_STREAM_SIZES, &sd->streamSizes) == SEC_E_OK)
  {
    sd->ioBufferSize = sd->streamSizes.cbHeader + sd->streamSizes.cbMaximumMessage + sd->streamSizes.cbTrailer;
    if((sd->ioBuffer = Mem::alloc(sd->ioBufferSize)) != NULL)
    {
      sd->socket = socket;
      return true;
    }
  }

  if(!FAILED(ssFirst))
  {
    Mem::free(sd->extraBuffer);
    sd->extraBuffer = NULL;
    safeFreeContextBuffer(&OutBuffers[0]);
    CWA(secur32, DeleteSecurityContext)(&sd->ch);
  }

  CWA(secur32, FreeCredentialsHandle)(&sd->sh);
  return false;
}

void SslSocket::_close(SOCKETDATA *sd)
{
  // SSL.
  //FIXME: IE   ,           (       SSL).

  // .
  Mem::free(sd->pendingData);
  Mem::free(sd->extraBuffer);
  Mem::free(sd->ioBuffer);
  CWA(secur32, DeleteSecurityContext)(&sd->ch);
  CWA(secur32, FreeCredentialsHandle)(&sd->sh);
}

bool SslSocket::_send(SOCKETDATA *sd, void *buf, int bufSize)
{
  LPBYTE pMessage    = (LPBYTE)sd->ioBuffer + sd->streamSizes.cbHeader;
  LPBYTE pTrailer    = pMessage + sd->streamSizes.cbMaximumMessage;
  DWORD dwDataToSend = (DWORD)bufSize;
  LPBYTE pPos        = (LPBYTE)buf;

  SecBufferDesc Message;
  SecBuffer     Buffers[4];

  while(dwDataToSend != 0)
  {
    Mem::_zero(sd->ioBuffer, sd->ioBufferSize);

    DWORD dwCurLen = min(dwDataToSend, sd->streamSizes.cbMaximumMessage);
    Mem::_copy(pMessage, pPos, dwCurLen);

    dwDataToSend -= dwCurLen;
    pPos += dwCurLen;

    // .
    Buffers[0].pvBuffer   = sd->ioBuffer;
    Buffers[0].cbBuffer   = sd->streamSizes.cbHeader;
    Buffers[0].BufferType = SECBUFFER_STREAM_HEADER;

    Buffers[1].pvBuffer   = pMessage;
    Buffers[1].cbBuffer   = dwCurLen;
    Buffers[1].BufferType = SECBUFFER_DATA;

    Buffers[2].pvBuffer   = pTrailer;
    Buffers[2].cbBuffer   = sd->streamSizes.cbTrailer;
    Buffers[2].BufferType = SECBUFFER_STREAM_TRAILER;

    Buffers[3].pvBuffer   = NULL;
    Buffers[3].cbBuffer   = 0;
    Buffers[3].BufferType = SECBUFFER_EMPTY;

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

    //.
    if(CWA(secur32, EncryptMessage)(&sd->ch, 0, &Message, 0) != SEC_E_OK)return false;

    //.
    for(BYTE i = 0; i < 3; i++)if(Buffers[i].cbBuffer > 0 && !WSocket::tcpSend(sd->socket, Buffers[i].pvBuffer, Buffers[i].cbBuffer))return false;
  }
  return true;
}

int SslSocket::_recv(SOCKETDATA *sd, void *buf, int bufSize, DWORD timeout, bool *extraAvalible)
{
  //      .
  if(sd->pendingDataSize > 0)
  {
    #if defined(WDEBUG0)
    WDEBUG0(WDDT_INFO, "Copying pending data.");
    #endif

    DWORD size = min(sd->pendingDataSize, (DWORD)bufSize);
    Mem::_copy(buf, sd->pendingData, size);
    sd->pendingDataSize -= size;

    if(sd->pendingDataSize == 0)
    {
      Mem::free(sd->pendingData);
      sd->pendingData = NULL;
    }
    else Mem::_copy(sd->pendingData, (LPBYTE)sd->pendingData + size, sd->pendingDataSize);
    *extraAvalible = (sd->extraBufferSize > 0 || sd->pendingDataSize > 0) ? true : false;
    return (int)size;
  }

  //  .
  SecBufferDesc Message;
  SecBuffer     Buffers[4];
  DWORD dwReaded = 0;
  int r = SOCKET_ERROR;

  for(;;)
  {
    DWORD dwCurLen;

    // .
    if(sd->extraBufferSize > 0)
    {
      #if defined(WDEBUG0)
      WDEBUG0(WDDT_INFO, "Copying extra data.");
      #endif

      dwCurLen = sd->extraBufferSize;
      Mem::_copy(sd->ioBuffer, sd->extraBuffer, sd->extraBufferSize);
      Mem::free(sd->extraBuffer);
      sd->extraBuffer = NULL;
      sd->extraBufferSize = 0;
    }
    else
    {
      dwCurLen = (DWORD)WSocket::tcpRecv(sd->socket, (LPBYTE)sd->ioBuffer + dwReaded, sd->ioBufferSize - dwReaded, timeout);
      if(dwCurLen == 0 || dwCurLen == (DWORD)SOCKET_ERROR){r = (int)dwCurLen; break;}
    }

    dwReaded += dwCurLen; //   .

    // .
    Buffers[0].pvBuffer   = sd->ioBuffer;
    Buffers[0].cbBuffer   = dwReaded;
    Buffers[0].BufferType = SECBUFFER_DATA;

    for(BYTE i = 1; i < 4; i++)
    {
      Buffers[i].pvBuffer   = NULL;
      Buffers[i].cbBuffer   = 0;
      Buffers[i].BufferType = SECBUFFER_EMPTY;
    }

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

    // .
    SECURITY_STATUS ss = CWA(secur32, DecryptMessage)(&sd->ch, &Message, 0, NULL);
    if(ss == SEC_E_INCOMPLETE_MESSAGE)
    {
#     if defined(WDEBUG0)
      WDEBUG0(WDDT_INFO, "SEC_E_INCOMPLETE_MESSAGE detected.");
#     endif
      continue;
    }

    //FIXME:  SEC_I_CONTEXT_EXPIRED (       _Close), SEC_I_RENEGOTIATE     .
    if(ss == SEC_E_OK)
    {
      //  .
      SecBuffer *pData  = NULL;
      SecBuffer *pExtra = NULL;
      for(BYTE i = 1; i < 4; i++)
      {
        if(pData == NULL && Buffers[i].BufferType == SECBUFFER_DATA && Buffers[i].cbBuffer > 0)pData = &Buffers[i];
        else if(pExtra == NULL && Buffers[i].BufferType == SECBUFFER_EXTRA && Buffers[i].cbBuffer > 0)pExtra = &Buffers[i];
      }

      //   pBuf.
      if(pData)
      {
        #if defined(WDEBUG1)
        WDEBUG1(WDDT_INFO, "SECBUFFER_DATA detected %u.", pData->cbBuffer);
        #endif
        if(pData->cbBuffer > (DWORD)bufSize)
        {
          dwCurLen = (DWORD)bufSize;
          sd->pendingDataSize = pData->cbBuffer - dwCurLen;
          if((sd->pendingData = Mem::copyEx((LPBYTE)pData->pvBuffer + dwCurLen, sd->pendingDataSize)) == NULL)break;
        }
        else dwCurLen = pData->cbBuffer;
        Mem::_copy(buf, pData->pvBuffer, dwCurLen);
      }

      //  .
      if(pExtra != NULL)
      {
        #if defined(WDEBUG1)
        WDEBUG1(WDDT_INFO, "SECBUFFER_EXTRA detected %u.", pExtra->cbBuffer);
        #endif

        if((sd->extraBuffer = Mem::copyEx((LPBYTE)sd->ioBuffer + dwReaded - pExtra->cbBuffer, pExtra->cbBuffer)) == NULL)break;
        sd->extraBufferSize = pExtra->cbBuffer;

        /*
          Sometimes an application will read data from the remote party, attempt to decrypt it by
          using DecryptMessage (Schannel), and discover that DecryptMessage (Schannel) succeeded
          but the output buffers are empty. This is normal behavior, and applications must be able
          to deal with it.

           .
        */
        if(pData == NULL)continue;
      }

      if(pData)r = (int)dwCurLen;
    }
    break;
  }

  *extraAvalible = (sd->extraBufferSize > 0 || sd->pendingDataSize > 0) ? true : false;
  return r;
}

bool SslSocket::_recvAll(SOCKETDATA *sd, void *buf, int bufSize, DWORD timeout, bool *extraAvalible)
{
  int len;
  do
  {
    if((len = _recv(sd, buf, bufSize, timeout, extraAvalible)) <= 0)return false;

    buf      = ((LPBYTE)buf) + len;
    bufSize -= len;
  }
  while(bufSize > 0);

  return bufSize == 0 ? true : false;
}
