#pragma once

#include <sspi.h>
#include <wincrypt.h>
#include <schnlsp.h>
#include <winsock2.h>
#include <wincrypt.h>
#include <security.h>
#include <assert.h>

// Status of SSLEngine
enum SSLEngineStatus
{
  BUFFER_OVERFLOW,
  BUFFER_UNDERFLOW,
  CLOSED,
  OK,
};

/**
 * Handshake status of SSLEngine.
 */
enum SSLEngineHandshakeStatus
{
  FINISHED,
  NEED_TASK,
  NEED_UNWRAP,
  NEED_WRAP,
  NOT_HANDSHAKING,
};

/**
 * Task result of SSLEngine.
 */
enum SSLEngineTaskStatus
{
  TASK_OK,
  TASK_CERTIFICATE_VERIFY_ERROR,
  TASK_ILLEGAL_STATE_ERROR
};

/**
 * Handshake trigger
 */
enum SSLEngineHandshakeTrigger
{
  HELLO_REQUEST,
  CLOSE_NOTIFY_RECEIVED,
  OUTBOUND_CLOSED
};

/**
 * SSLEngine実行結果を表すクラス
 */
class SSLEngineResult
{
private:
  SSLEngineStatus status;
  SSLEngineHandshakeStatus handshakeStatus;

  /**
   * Bytes consumed.
   */
  int consumed;

  /**
   * Bytes produced.
   */
  int produced;

  /**
   * Result of SSPI
   */
  HRESULT result;

public:
  SSLEngineResult()
    : status(OK),
      handshakeStatus(NOT_HANDSHAKING),
      consumed(0),
      produced(0),
      result(S_OK)
  {
  }

  inline void setValues(SSLEngineStatus status, SSLEngineHandshakeStatus handshakeStatus, int bytesConsumed, int bytesProduced)
  {
    this->status = status;
    this->handshakeStatus = handshakeStatus;
    this->consumed = bytesConsumed;
    this->produced = bytesProduced;
    this->result = S_OK;
  }

  inline void setError(HRESULT errorResult)
  {
    assert(FAILED(errorResult));
    this->result = errorResult;

    this->status = OK;
    this->handshakeStatus = NOT_HANDSHAKING;
    this->consumed = 0;
    this->produced = 0;

    assert(false);
  }

  inline HRESULT getResult() const
  {
    return this->result;
  }

  inline SSLEngineStatus getStatus() const
  {
    return this->status;
  }

  inline SSLEngineHandshakeStatus getHandshakeStatus() const
  {
    return this->handshakeStatus;
  }

  inline int bytesConsumed() const
  {
    return this->consumed;
  }

  inline int bytesProduced() const
  {
    return this->produced;
  }
};

/**
 * SSLEngine
 */
class SSLEngine
{
private:
  CRITICAL_SECTION          criticalSection;

  const _TCHAR*             peerHost;
  CredHandle*               phCredential;
  CtxtHandle*               phContext;

  static PSecurityFunctionTable    pFunctions;

  /**
   * クライアントモードで動作する場合true
   */
  bool                      clientMode;


  /**
   * 出力側が閉じているか
   */
  bool                      outboundDone;

  /**
   * 入力側が閉じているか
   */
  bool                      inboundDone;

  /**
   * ハンドシェークステータス
   */
  SSLEngineHandshakeStatus  handshakeStatus;
  
  /**
   * 最後に実施したハンドシェーク動作のトリガー
   */
  SSLEngineHandshakeTrigger lastHandshakeTrigger;

  /**
   * 直前のInitializeSecurityContext()/AcceptSecurityContext() 呼び出し結果
   */
  HRESULT                   securityStatus;

  /**
   * InitializeSecurityContext()/AcceptSecurityContext()の出力バッファ
   */
  SecBuffer securityContextOutput[1];
  
  /**
   * wrap()の結果
   */
  SSLEngineResult wrapResult;

  /**
   * unwrap()の結果
   */
  SSLEngineResult unwrapResult;

  /**
   * 有効になっているプロトコル
   */
  DWORD enabledProtocols;

  /**
   * 使用されているプロトコルのバージョン
   */
  unsigned short protocolVersion;

  /**
   * セッションIDの長さ
   */
  unsigned char sessionIDLength;

  /**
   * セッションID
   */
  unsigned char* sessionID;

  /**
   * 使用されている Cipher suiteをあらわすコード
   */
  unsigned short cipherSuiteCode;

  /**
   * 自身の証明書
   */
  PCCERT_CONTEXT pLocalCertificate;

  /**
   * 自身の証明書チェーン
   */
  PCCERT_CHAIN_CONTEXT pLocalCertificateChain;

  /**
   * 相手側の証明書
   */
  PCCERT_CONTEXT pPeerCertificate;

  /**
   * 相手側の証明書チェーン
   */
  PCCERT_CHAIN_CONTEXT pPeerCertificateChain;

  /**
   * クライアント認証レベル
   *
   * 0 - 認証なし
   * 1 - 認証を要求する (wantClientAuth)
   * 2 - 認証が必須     (needClientAuth)
   */
  int clientAuthLevel;

  /**
   * ストリームバッファサイズ
   */
  SecPkgContext_StreamSizes streamSizes;

  /**
   * 相手局の証明書取得処理が完了したか
   */
  bool peerCertificateChainLoaded;

  /**
   * 委譲タスク数
   */
  int delegatedTaskCount;

  //-----------------------------------------------------------
  // メンバ関数
  //-----------------------------------------------------------

  /**
   * ロックする
   */
  inline void acquireLock()
  {
    EnterCriticalSection(&this->criticalSection);
  }

  /**
   * ロックを解除する
   */
  inline void releaseLock()
  {
    LeaveCriticalSection(&this->criticalSection);
  }

  /**
   * close_notifyアラートを生成する
   */
  void wrapCloseNotify();

  inline void setHandshakeStatus(SSLEngineHandshakeStatus stat)
  {
    this->handshakeStatus = stat;
  }

  /**
   * InitializeSecurityContext()の初回呼び出しを行う
   */
  void firstInitializeSecurityContext(void* destbuff, int destlen);

  /**
   * AcceptSecurityContext()の初回呼び出しを行う
   */
  void firstAcceptSecurityContext(void* srcbuff, int srclen);

  /**
   * 2回目以降のInitializeSecurityContext()/AcceptSecurityContext()の呼び出し処理
   */
  void wrapSecurityContext(void* destbuff, int destlen);
  void unwrapSecurityContext(void* srcbuff, int srclen);

  /**
   * ハンドシェーク状態を更新する
   */
  bool updateHandshakeStatus();

  /**
   * ハンドシェイク動作を実施した原因を返す
   */
  inline SSLEngineHandshakeTrigger getLastHandshakeTrigger() const
  {
    return this->lastHandshakeTrigger;
  }

  inline void setLastHandshakeTrigger(SSLEngineHandshakeTrigger trigger)
  {
    this->lastHandshakeTrigger = trigger;
  }

  inline void setOutboundDone(bool done)
  {
    this->outboundDone = done;
  }

  inline void setInboundDone(bool done)
  {
    this->inboundDone = done;
  }

  /**
   * アプリケーションデータを暗号化する
   */
  void encryptApplicationData(void* srcbuff, int srcdata, void* destbuff, int destlen);

  /**
   * ネットワークデータを復号化する
   */
  void decryptPacketData(void* srcbuff, int srcdata, void* destbuff, int destlen);

  /**
   * 指定されたデータがServer helloであった場合、protocol versionとcipher suiteを取得する
   */
  bool parseServerHello(void* buff, int len);

  /**
   * 相手側から受信した close_notify を処理する
   */
  void onCloseNotifyReceived(int bytesConsumed);

public:
  SSLEngine(const _TCHAR* host);
  ~SSLEngine();

  inline const _TCHAR* getPeerHost() const
  {
    return this->peerHost;
  }

  /**
   * アプリケーションデータを暗号化する
   */
  const SSLEngineResult& wrap(void* srcbuff, int srclen, void* destbuff, int destlen);

  /**
   * 暗号化されたデータを復号する。
   */
  const SSLEngineResult& unwrap(void* srcbuff, int srclen, void* destbuff, int destlen);

  inline SSLEngineHandshakeStatus getHandshakeStatus()
  {
    SSLEngineHandshakeStatus hs = this->handshakeStatus;
    if (FINISHED == hs)
    {
      // FINISHEDからNOT_HANDSHAKINGに移行させる
      this->handshakeStatus = NOT_HANDSHAKING;
    }
    return hs;
  }

  /**
   * クライアントとしてハンドシェークを行うかどうかを設定する
   *
   * @param  cm
   */
  inline void setUseClientMode(bool cm)
  {
    this->clientMode = cm;
  }

  /**
   * クライアントとしてハンドシェークを行うかどうかを返す
   *
   * @return
   */
  inline bool getUseClientMode()
  {
    return this->clientMode;
  }

  /**
   * ハンドシェークを開始する
   */
  HRESULT beginHandshake();

  /**
   * 委譲タスクが存在するかを返す
   */
  bool hasDelegatedTask();

  /**
   * 委譲タスクが存在する場合、そのタスクを実行する
   */
  SSLEngineTaskStatus executeDelegatedTask();

  /**
   * 相手側証明書を取得する
   */
  bool loadPeerCertificateChain();

  /**
   * 相手側の証明書チェーンを返す
   */
  inline PCCERT_CHAIN_CONTEXT getPeerCertificateChain() const
  {
    return this->pPeerCertificateChain;
  }

  /**
   * 自分の証明書を返す
   */
  inline PCCERT_CONTEXT  getLocalCertificate() const
  {
    return this->pLocalCertificate;
  }

  /**
   * 自分の証明書チェーンを返す
   */
  PCCERT_CHAIN_CONTEXT getLocalCertificateChain();

  /**
   * 出力をクローズする
   */
  void closeOutbound();

  /**
   * 入力をクローズする
   */
  bool closeInbound();

  inline bool isOutboundDone() const
  {
    return this->outboundDone;
  }

  inline bool isInboundDone() const
  {
    return this->inboundDone;
  }

  /**
   * 利用可能なプロトコルを設定する
   */
  inline void setEnabledProtocols(DWORD protocols)
  {
    this->enabledProtocols = protocols;
  }

  inline DWORD getEnabledProtocols() const
  {
    return this->enabledProtocols;
  }

  inline unsigned short getProtocolVersion() const
  {
    return this->protocolVersion;
  }

  inline unsigned short getCipherSuiteCode() const
  {
    return this->cipherSuiteCode;
  }

  inline unsigned char getSessionIDLength() const
  {
    return this->sessionIDLength;
  }

  inline unsigned char* getSessionID() const
  {
    return this->sessionID;
  }

  /**
   * 自局側の証明書を設定する
   */
  void setLocalCertificate(PCCERT_CONTEXT localCert);

  inline void setWantClientAuth(bool auth)
  {
    this->clientAuthLevel = auth ? 1 : 0;
  }

  inline bool getWantClientAuth() const
  {
    return this->clientAuthLevel == 1;
  }

  inline void setNeedClientAuth(bool auth)
  {
    this->clientAuthLevel = auth ? 2 : 0;
  }

  inline bool getNeedClientAuth() const
  {
    return this->clientAuthLevel == 2;
  }

  /**
   * パケットデータサイズを返す。
   * 不明な場合は 0 を返す。
   */
  inline unsigned long getPacketBufferSize() const
  {
    unsigned long message = this->streamSizes.cbMaximumMessage;
    if (0 < message && message < 16384)
    {
      // http://support.microsoft.com/?scid=kb%3Ben-us%3B300562&x=13&y=13
      // によると、cbMaximumMessageの値は、SSLの仕様である16384よりも5バイト小さな値16379を
      // 返してくるプラットフォームがある。
      message = 16384;
    }
    return this->streamSizes.cbHeader + message + this->streamSizes.cbTrailer;
  }

  /**
   * アプリケーションデータバッファサイズを返す
   */
  inline unsigned long getApplicationBufferSize() const
  {
    unsigned long message = this->streamSizes.cbMaximumMessage;
    if (0 < message && message < 16384)
    {
      // http://support.microsoft.com/?scid=kb%3Ben-us%3B300562&x=13&y=13
      // によると、cbMaximumMessageの値は、SSLの仕様である16384よりも5バイト小さな値16379を
      // 返してくるプラットフォームがある。
      message = 16384;
    }
    return message;
  }
};