/**
 * @RDE_HEADER@
 * $Id$
 * $HeadURL$
 * $Author$
 */

#include "pws_vnc_handler.h"
#include "pws_des_encryption.h"

namespace ProjectionWorkstation
{


VncHandler::VncHandler(const QHostAddress &host, quint16 port, QObject *parent)
  : QObject(parent)
  , m_logger(LtsProjectionWorkstation::getLogger("VncHandler"))
  , m_nextType(RfbProtocol::STC::RFB_VERSION)
  , m_socket(new QTcpSocket(this))
  , m_host(host)
  , m_port(port)
{
  connect(m_socket, &QTcpSocket::readyRead, this, &VncHandler::onReadyRead);
  connect(m_socket, QOverload<QAbstractSocket::SocketError>::of(&QTcpSocket::error), this, &VncHandler::error);
  connect(m_socket, &QTcpSocket::disconnected, this, &VncHandler::slotOnDisconnected);
}

void VncHandler::connectToHost()
{
  m_socket->connectToHost(m_host, m_port);
}

void VncHandler::connectToHost(const QHostAddress &host, quint16 port)
{
  m_socket->connectToHost(host, port);
}

void VncHandler::onReadyRead()
{
  m_buffer += m_socket->readAll();

  QDataStream stream(m_buffer);
  stream.setByteOrder(QDataStream::BigEndian);
  qint64 readPos = 0;
  while (handleData(stream))
    readPos = stream.device()->pos();

  // TODO: squeeze m_buffer only when needed (e.g is larger than x MB)
  if (readPos > 0)
    m_buffer = m_buffer.mid(readPos);
}

bool VncHandler::handleData(QDataStream &stream)
{
  switch (m_nextType)
  {
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_VERSION:
  {
    RfbProtocol::STC::RfbVersion version;
    stream.startTransaction();
    stream >> version;
    if (!stream.commitTransaction())
      return false;
    if (handleVersion(version))
    {
      setNextExpectedMessageType(RfbProtocol::STC::RFB_SUP_SECURITY_TYPES);
      DEBUG_LOG(m_logger) << "RfbVersion handled Successfully!";
      return true;
    }
    ERROR_LOG(m_logger) << "RfbVersion handshake Error!";
    break;
  }
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_SUP_SECURITY_TYPES:
  {
    RfbProtocol::STC::RfbSupportedSecurityTypes supTypes;
    stream.startTransaction();
    stream >> supTypes;
    if (!stream.commitTransaction())
      return false;
    if (handleSupportedSecurityTypes(supTypes))
    {
      setNextExpectedMessageType(RfbProtocol::STC::RFB_VNC_AUTHENTICATION);
      DEBUG_LOG(m_logger) << "RfbSecurityType handled Successfully!";
      return true;
    }
    ERROR_LOG(m_logger) << "RfbSecurityType handshake Error!";
    break;
  }
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_VNC_AUTHENTICATION:
  {
    RfbProtocol::STC::RfbVncAuthentication auth;
    stream.startTransaction();
    stream >> auth;
    if (!stream.commitTransaction())
      return false;
    if (handleVncAuthentication(auth))
    {
      setNextExpectedMessageType(RfbProtocol::STC::RFB_SECURITY_RESULT);
      DEBUG_LOG(m_logger) << "RfbVncAuthentication handled Successfully!";
      return true;
    }
    ERROR_LOG(m_logger) << "RfbVncAuthentication Error!";
    break;
  }
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_SECURITY_RESULT:
  {
    RfbProtocol::STC::RfbSecurityResult result;
    stream.startTransaction();
    stream >> result;
    if (!stream.commitTransaction())
      return false;
    if (handleSecurityResult(result))
    {
      setNextExpectedMessageType(RfbProtocol::STC::RFB_SERVER_INIT);
      DEBUG_LOG(m_logger) << "RfbSecurityResult handled Successfully!";
      return true;
    }
    ERROR_LOG(m_logger) << "RfbSecurityResult Error!";
    break;
  }
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_SERVER_INIT:
  {
    RfbProtocol::STC::RfbServerInit init;
    stream.startTransaction();
    stream >> init;
    if (!stream.commitTransaction())
      return false;
    if (handleServerInit(init))
    {
      QByteArray data = RfbProtocol::CTS::RfbSetEncodings({qint32(RfbProtocol::HEXTILE_ENCODING),
                                                           qint32(RfbProtocol::RAW_ENCODING),
                                                           qint32(RfbProtocol::LAST_RECT_PSEUDO_ENCODING),
                                                           qint32(-247),  // Compression Level Pseudo-encoding (-247 high .. -255 low)
                                                           qint32(-23),   // JPEG Quality Level Pseudo-encoding (-23 high quality .. -32 low quality)
                                                          });

      m_socket->write(data);

      emit connected();
      m_socket->write(RfbProtocol::CTS::RfbFrameBufferUpdateRequest(false, 0, 0, 1920, 1200));
      setNextExpectedMessageType(RfbProtocol::STC::RFB_FRAMEBUFFER_UPDATE);
      DEBUG_LOG(m_logger) << "RfbServerInit handled Successfully!";
    }
    return true;
  }
  // -------------------------------------------------------------------------------------------
  case RfbProtocol::STC::RFB_FRAMEBUFFER_UPDATE:
  {
    // The handling of 'next expected type is flawed' - the order and next type is only defined during setup,
    // not during runtime - here we can get any Server to Client Message defined in the rfb protocol. So we
    // have to take a look at the protocol byte and handle accordingly
    // by default we must only be able to decode 0 (Framebuffer update), 1 (SetColourMapEntries), 2 (Bell) and 3 (ServerCutText)
    // also using QDataStream here only works by accident... but we won't change it now (never touch a running system)
    stream.startTransaction();
    quint8 msgType;
    stream >> msgType;
    DEBUG_LOG(m_logger) << "got message from server: " << QString::number(msgType, 16);
    switch (msgType)
    {
    case RfbProtocol::STC::RFB_FRAMEBUFFER_UPDATE:
      return handleFrameBufferUpdate(stream);
    case RfbProtocol::STC::RFB_SET_COLORMAP:
      return handleSetColormap(stream);
    case RfbProtocol::STC::RFB_RING_BELL:
      return handleRingBell(stream);
    case RfbProtocol::STC::RFB_SERVER_CUTTEXT:
      return handleServerCuttext(stream);
    default:
      WARNING_LOG(m_logger) << "Unknown message type from server (" << QString::number(msgType, 16) << ") - reconnecting";
      stream.abortTransaction();
      slotOnDisconnected();
      return false;
    }
    return true;
  }
  default:
    break;
  }
  Q_ASSERT(false);
  stream.abortTransaction();
  slotOnDisconnected();
  return false;
}

bool VncHandler::handleFrameBufferUpdate(QDataStream &stream)
{
  RfbProtocol::STC::RfbFramebufferUpdate update;
  stream.startTransaction();
  stream >> update;
  bool bRet = stream.commitTransaction();
  if (bRet)
  {
    if (!update.rectangles().isEmpty())
      emit frameBufferUpdateReceived(update.rectangles());
    m_socket->write(RfbProtocol::CTS::RfbFrameBufferUpdateRequest(true, 0, 0, 1920, 1200));
  }
  return bRet;
}

bool VncHandler::handleSetColormap(QDataStream &stream)
{
  RfbProtocol::STC::RfbSetColormap setColormap;
  stream.startTransaction();
  stream >> setColormap;
  return stream.commitTransaction();
}

bool VncHandler::handleRingBell(QDataStream &stream)
{
  RfbProtocol::STC::RfbRingBell cutText;
  stream.startTransaction();
  stream >> cutText;
  bool bRet = stream.commitTransaction();
    emit beepReceived();
  return bRet;
}

bool VncHandler::handleServerCuttext(QDataStream &stream)
{
  RfbProtocol::STC::RfbServerCutText cutText;
  stream.startTransaction();
  stream >> cutText;
  bool bRet = stream.commitTransaction();
  if (bRet)
    emit cutTextReceived(QString::fromLatin1(cutText.text()));
  return bRet;
}

void VncHandler::slotOnDisconnected()
{
  resetHandler();
  emit disconnected();
}

void VncHandler::writeData(const QByteArray &data)
{
  m_socket->write(data);
  m_socket->flush();
}

void VncHandler::error()
{
  ERROR_LOG(m_logger) << "SocketError " << m_socket->errorString();
}

void VncHandler::resetHandler()
{
  m_buffer.clear();
  setNextExpectedMessageType(RfbProtocol::STC::RFB_VERSION);
  m_serverInit = RfbProtocol::STC::RfbServerInit();
}

// Data Handlers

bool VncHandler::handleVersion(const RfbProtocol::STC::RfbVersion &version)
{
  if (version.majorVersion() == 3 && version.minorVersion() == 8)
  {
    m_socket->write(RfbProtocol::CTS::RfbVersion(version.majorVersion(),
                                                 version.minorVersion()));
    return true;
  }
  ERROR_LOG(m_logger) << "Unsupported Version: " << version.majorVersion() << "." << version.minorVersion();
  return false;
}

bool VncHandler::handleSupportedSecurityTypes(const RfbProtocol::STC::RfbSupportedSecurityTypes &supportedTypes)
{
  if (supportedTypes.types().contains(2) && !supportedTypes.types().contains(0))
  {
    m_socket->write(RfbProtocol::CTS::RfbSecurityType(2));
    return true;
  }
  QString typesString;
  for (quint8 type : supportedTypes.types())
    typesString.append(QString::number(type)).append(QLatin1Char(' '));
  if (typesString.isEmpty()) typesString = QStringLiteral("EMPTY");
  ERROR_LOG(m_logger) << "Unsupported Security Types: " << typesString;
  return false;
}

bool VncHandler::handleVncAuthentication(const RfbProtocol::STC::RfbVncAuthentication &authentication)
{
  quint64 key = convertPassword("ctc11de");
  quint64 blockA = 0, blockB = 0;
  DES des(key);
  blockA = des.encrypt(authentication.challengeBlockA());
  blockB = des.encrypt(authentication.challengeBlockB());
  m_socket->write(RfbProtocol::CTS::RfbVncAuthentication(blockA, blockB));
  return true;
}

// VNC Servers require the password with reversed bytes
// and a length of 8bytes (truncated or padded if longer/smaller)
quint64 VncHandler::convertPassword(const char *password)
{
  quint64 result = 0;
  const size_t len = strlen(password);
  for (quint8 count = 0; count < 8; ++count)
  {
    quint8 data;
    result <<= 8;
    if (count < len)
      data = reverseByte(quint8(password[count]));
    else
      data = 0;
    result = (result | data);
  }
  return result;
}

quint8 VncHandler::reverseByte(quint8 b)
{
  b = (b & 0xF0) >> 4 | (b & 0x0F) << 4;
  b = (b & 0xCC) >> 2 | (b & 0x33) << 2;
  b = (b & 0xAA) >> 1 | (b & 0x55) << 1;
  return b;
}

bool VncHandler::handleSecurityResult(const RfbProtocol::STC::RfbSecurityResult &result)
{
  if (result.resultCode() == 0)
  {
    m_socket->write(RfbProtocol::CTS::RfbClientInit(true));
    return true;
  }
  else
  {
    ERROR_LOG(m_logger) << result.reason();
    return false;
  }
}

bool VncHandler::handleServerInit(const RfbProtocol::STC::RfbServerInit &serverInit)
{
  m_serverInit = serverInit;
  return true;
}

} // ProjectionWorkstation