//
//  raspVNC2: VNC module (new version)
//      xkozima@myu.ac.jp

#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <stdio.h>
#include <string.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <rpc/des_crypt.h>
#include "raspVNC2.h"

#define DEF_SERVER_ADDR "127.0.0.1"
#define DEF_SERVER_PORT 5901

//  コンストラクタ
VNC::VNC () {
    //  set default values
    addr = DEF_SERVER_ADDR;
    port = DEF_SERVER_PORT;
    sock = 0;
    desKey  = NULL;
    title   = NULL;
    drawCB  = NULL;
    serverX = 0;
    serverY = 0;
    buffer  = NULL;
    bufferSize = 0;
    redrawFlag = false;
    touchCB = NULL;
    touchLastM = VNC_TOUCH_NONE;
}

//  サーバを設定（必要に応じて init() より前に呼び出す）
void VNC::presetServer (const char *serverAddr, int serverPort) {
    addr = strdup(serverAddr);          //  VNC::quit() で解放
    port = serverPort;
}
//  コールバックを設定（必要に応じて init() より前に呼び出す）
void VNC::presetDrawCB (void (*cb)(unsigned char *buf, int x, int y, int w, int h)) {
    drawCB = cb;
}
void VNC::presetTouchCB (void (*cb)(int *x, int *y, int *mode)) {
    touchCB = cb;
}
//  パスワードを設定（必要に応じて init() より前に呼び出す）
void VNC::presetPassword (const char *password) {
    char pass8[8];
    //  パスワードを８文字にする
    for (int i = 0; i < 8; i++) {
        if (i < strlen(password))
            pass8[i] = password[i];
        else
            pass8[i] = 0;
    }
    //  ビットの順序を入れ替え（なぜ？）
    for (int i = 0; i < 8; i++) {
        char ch = pass8[i];
        pass8[i] = ((ch & 0x80)? 0x01: 0)
                 | ((ch & 0x40)? 0x02: 0)
                 | ((ch & 0x20)? 0x04: 0)
                 | ((ch & 0x10)? 0x08: 0)
                 | ((ch & 0x08)? 0x10: 0)
                 | ((ch & 0x04)? 0x20: 0)
                 | ((ch & 0x02)? 0x40: 0)
                 | ((ch & 0x01)? 0x80: 0);
    }
    //  LSB にパリティ（奇数）を入れる
    for (int i = 0; i < 8; i++) {
        int parity = 0;
        for (int j = 0; j < 8; j++)
            if (pass8[i] & (1 << j)) parity++;
        if ((parity & 1) == 0)
            pass8[i] |= 0x01;
    }
    //  desKey に記憶（セキュリティは大丈夫か？）
    desKey = strdup(pass8);
}

//  VNC 初期化（サーバとの接続）
void VNC::init (int w, int h, int fmt) {
    int res;
    //  サーバに接続
    struct sockaddr_in address;
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = inet_addr(addr);
    address.sin_port = htons(port);
    sock = socket(AF_INET, SOCK_STREAM, 0);
    res = connect(sock, (struct sockaddr *) &address, sizeof(address));
    if (res == -1) {
        printf("VNC (error): cannot connect to server (VNC)\n");
        exit(1);
    }
    //  Handshake: ハンドシェイク
    unsigned char serverVer[12];
    secureRead(sock, serverVer, 12);    //  たぶん "RFB 003.008"
    unsigned char clientVer[13] = "RFB 003.003\n";
    secureWrite(sock, clientVer, 12);   //  こちらは Ver 3.3
    //  Authentification: セキュリティ認証
    unsigned char auth[4];
    secureRead(sock, auth, 4);
    if (auth[0] == 0 && auth[1] == 0 && auth[2] == 0) {
        if (auth[3] == 0) {
            //  ハンドシェーク失敗（auth: 0）
            printf("VNC (error): cannot complete handshake (VNC)\n");
            exit(1);
        }
        else if (auth[3] == 1) {
            //  認証不要（auth: 1）
            printf("VNC (info): no authentificationi required\n");
        }
        else if (auth[3] == 2) {
            //  パスワード認証（auth: 2）
            if (desKey == NULL) {
                printf("VNC (error): password has not been set (VNC)\n");
                exit(1);
            }
            //  チャレンジ受信・暗号化・レスポンス送信
            unsigned char chaRes[16];
            secureRead(sock, chaRes, 16);
            ecb_crypt(desKey, (char *) chaRes, 16, DES_ENCRYPT);
            secureWrite(sock, chaRes, 16);
            //  結果をチェック（0: 成功, non-zero: 失敗）
            unsigned char result[4];
            secureRead(sock, result, 4);
            if (result[0] | result[1] | result[2] | result[3]) {
                printf("VNC (error): authentification failed (VNC)\n");
                exit(1);
            }
            printf("VNC (info): VNC authentificationi successful\n");
        }
        else {
            //  0/1/2 以外は通信エラー
            printf("VNC (error): communication error (VNC)\n");
            exit(1);
        }    
    }
    else {
        //  通信エラー
        printf("VNC (error): communication error (VNC)\n");
        exit(1);
    }    
    //  ClientInit: クライアントからサーバへ合図（１バイトの１）
    unsigned char clientInit[] = {1};
    secureWrite(sock, clientInit, 1);
    //  ServerInit: サーバからクライアントへ情報（幅/高/形式/タイトル）
    unsigned char strWh[4], strFmt[16], strLen[4];
    secureRead(sock, strWh, 4);
    serverWidth  = strWh[0] * 256 + strWh[1];
    serverHeight = strWh[2] * 256 + strWh[3];
    secureRead(sock, strFmt, 16);   //  RGB565 を想定（チェック省略）
    secureRead(sock, strLen, 4);
    printf("VNC (info): server PixelFormat is %dx%d@%dbpp\n", 
           serverWidth, serverHeight, strFmt[0] );
    int length = strLen[2] * 256 + strLen[3];
    title = (char *) malloc(length + 1);
    secureRead(sock, (unsigned char *) title, length);
    title[length] = '\0';
    //  SetPixelFormat: ピクセル形式の指定
    if (fmt == VNC_FMT_RGB16) {
        unsigned char setPixFmt[] = {
            0x00, 0x00, 0x00, 0x00,             //  Command pad3
            0x10, 0x10, 0x00, 0x01,             //  16bpp/16dep/LE/trueC
            0x00, 0x1f, 0x00, 0x3f, 0x00, 0x1f, //  RGB max
            0x0b, 0x05, 0x00,                   //  RGB shift
            0x00, 0x00, 0x00 };                 //  pad3
        secureWrite(sock, setPixFmt, 20);
        printf("VNC (info): set PixelFormat to %dbpp/%ddepth\n", 
               setPixFmt[4], setPixFmt[5] );
        bytesPP = 2;
    }
    else if (fmt == VNC_FMT_RGB24) {
        unsigned char setPixFmt[] = {
            0x00, 0x00, 0x00, 0x00,             //  Command pad3
            0x20, 0x18, 0x00, 0x01,             //  32bpp/24dep/LE/trueC
            0x00, 0xff, 0x00, 0xff, 0x00, 0xff, //  RGB max
            0x10, 0x08, 0x00,                   //  RGB shift
            0x00, 0x00, 0x00 };                 //  pad3
        secureWrite(sock, setPixFmt, 20);
        printf("VNC (info): set PixelFormat to %dbpp/%ddepth\n", 
               setPixFmt[4], setPixFmt[5] );
        bytesPP = 4;
    }
    else {
        printf("VNC (error): unsupported format\n");
        exit(1);
    }
    //  SetEncodings: エンコーディング形式の指定（RAW のみ）
    unsigned char setEnc[] = {2, 0, 0, 1, 0, 0, 0, 0};
    secureWrite(sock, setEnc, 8);
    printf("VNC (info): set Encodings to RAW\n");
    //  最初の FramebufferUpdateRequest（全画面）
    //      ROI = (0, 0, width, height)
    serverX = 0;  serverY = 0;
    width   = w;  height  = h;
    unsigned char fbReq[10] = {0x03, 0x00};
    fbReq[2] = (serverX >> 8) & 0xff;  fbReq[3] = serverX & 0xff;
    fbReq[4] = (serverY >> 8) & 0xff;  fbReq[5] = serverY & 0xff;
    fbReq[6] = (width   >> 8) & 0xff;  fbReq[7] = width   & 0xff;
    fbReq[8] = (height  >> 8) & 0xff;  fbReq[9] = height  & 0xff;
    secureWrite(sock, fbReq, 10);
}
//  VNC 終了（サーバ切断・メモリ解放）
void VNC::quit () {
    //  サーバから切断
    close(sock);
    //  メモリの解放
    if (desKey) free(desKey);
    if (title) free(title);
    if (buffer) free(buffer);
    bufferSize = 0;
}
//  タイトル取得
char *VNC::getTitle () {
    return title;
}

//  画面更新
//  （サーバから更新情報を読み出し，drawCB() を適用して画面描画）
//  （touchCB() でタッチ情報を取得し，必要に応じてサーバにイベントを送信）
void VNC::update () {
    //  FramebufferUpdate の処理
    int bytesAvailable;
    ioctl(sock, FIONREAD, &bytesAvailable);
    if (bytesAvailable) {
        unsigned char header[4];
        secureRead(sock, header, 4);
        if (header[0] == 0) {
            //  フレーム情報を読み出す
            int num = header[2] * 256 + header[3];
            for (int i = 0; i < num; i++) {
                unsigned char xy[4], wh[4], enc[4];
                secureRead(sock, xy,  4);
                secureRead(sock, wh,  4);
                secureRead(sock, enc, 4);
                int x = (xy[0] << 8) + xy[1];
                int y = (xy[2] << 8) + xy[3];
                int w = (wh[0] << 8) + wh[1];
                int h = (wh[2] << 8) + wh[3];
                int e = enc[3];
                //  バッファ確保
                int bytes = w * h * bytesPP;
                if (buffer == NULL) {
                    buffer = (unsigned char *) malloc(bytes);
                    bufferSize = bytes;
                }
                else if (bytes > bufferSize) {
                    free(buffer);
                    buffer = (unsigned char *) malloc(bytes);
                    bufferSize = bytes;
                }
                //  画像を読み出す
                secureRead(sock, buffer, bytes);
                //  コールバックで処理する
                if (drawCB)
                    drawCB(buffer, x - serverX, y - serverY, w, h);
            }
            //  FramebufferUpdateRequest（つぎの画面更新に向けて）
            //      ROI = (serverX, serverY, width, height)
            unsigned char fbReq[10] = {0x03};
            fbReq[1] = (redrawFlag)? 0x00: 0x01;  //  スクロール直後は
            redrawFlag = false;                   //  一回だけ全画面更新
            fbReq[2] = (serverX >> 8) & 0xff;  fbReq[3] = serverX & 0xff;
            fbReq[4] = (serverY >> 8) & 0xff;  fbReq[5] = serverY & 0xff;
            fbReq[6] = (width   >> 8) & 0xff;  fbReq[7] = width   & 0xff;
            fbReq[8] = (height  >> 8) & 0xff;  fbReq[9] = height  & 0xff;
            secureWrite(sock, fbReq, 10);
        }
        else if (header[0] == 1) {
            //  SetColorMapEntries (not implemented)
        }
        else if (header[0] == 2) {
            //  Bell (not implemented)
        }
        else if (header[0] == 3) {
            //  ServerCutText (not implemented)
        }
    }
    //  タッチ処理
    if (touchCB) {
        //  タッチ情報を取得
        int touchX, touchY, mode;
        touchCB(&touchX, &touchY, &mode);
        //  タッチダウン・ドラッグをサーバーへ送信
        if (VNC_TOUCH_1 <= mode && mode <= VNC_TOUCH_5) {
            //  ポインタ位置とボタン状態を送信・記録
            sendPointerEvent(serverX + touchX, serverY + touchY, mode);
            touchLastX = touchX;
            touchLastY = touchY;
            touchLastM = mode;
            usleep(5000);
        }
        else if (mode == VNC_TOUCH_SCR) {
            //  ドラッグ開始位置を記録
            if (touchLastM != VNC_TOUCH_SCR) {
                touchScrX = touchX;
                touchScrY = touchY;
            }
            //  ポインタのみ動かす
            sendPointerEvent(serverX + touchX, serverY + touchY, 0);
            touchLastX = touchX;
            touchLastY = touchY;
            touchLastM = mode;
            usleep(5000);
        }
        else if (touchLastM != VNC_TOUCH_NONE) {
            //  画面スクロールの処理
            if (touchLastM == VNC_TOUCH_SCR) {
                //  ドラッグ分だけ serverX/Y を移動
                serverX -= (touchLastX - touchScrX);
                serverY -= (touchLastY - touchScrY);
                if (serverX < 0)
                    serverX = 0;
                else if (serverX + width > serverWidth)
                    serverX = serverWidth - width;
                if (serverY < 0)
                    serverY = 0;
                else if (serverY + height > serverHeight)
                    serverY = serverHeight - height;
                redrawFlag = true;
            }
            //  タッチアップしたのでその位置を送信
            sendPointerEvent(serverX + touchLastX, serverY + touchLastY, 0);
            touchLastM = VNC_TOUCH_NONE;
            usleep(5000);
        }
        else {
            //  タッチなし（10ms 休み）
            usleep(10000);
        }
    }
    else {
        //  コールバック未登録なので 10ms 待ち
        usleep(1000);
    }
}

//  キーイベントの送信（keysym: X-Window の keysym）
void VNC::sendKeyPressed(unsigned int keysym) {
    unsigned char keyEv[8];
    keyEv[0] = 4;   //  Command
    keyEv[1] = 1;   //  1:keyDown, 0:keyUp
    keyEv[4] = (keysym >> 24) & 0xff;
    keyEv[5] = (keysym >> 16) & 0xff;
    keyEv[6] = (keysym >>  8) & 0xff;
    keyEv[7] = (keysym      ) & 0xff;
    secureWrite(sock, keyEv, 8);
}
void VNC::sendKeyReleassed(unsigned int keysym) {
    unsigned char keyEv[8];
    keyEv[0] = 4;   //  Command
    keyEv[1] = 0;   //  1:keyDown, 0:keyUp
    keyEv[4] = (keysym >> 24) & 0xff;
    keyEv[5] = (keysym >> 16) & 0xff;
    keyEv[6] = (keysym >>  8) & 0xff;
    keyEv[7] = (keysym      ) & 0xff;
    secureWrite(sock, keyEv, 8);
}
//  ポインタイベントの送信（button: 左=1, 中=2, 右=3, 上=4, 下=5）
void VNC::sendPointerEvent(int x, int y, int buttons) {
    unsigned char ptrEv[6];
    ptrEv[0] = 5;
    ptrEv[1] = (1 << buttons) >> 1;     //  左1,中2,右4,上8,下16
    ptrEv[2] = (x >> 8) & 0xff;
    ptrEv[3] = (x     ) & 0xff;
    ptrEv[4] = (y >> 8) & 0xff;
    ptrEv[5] = (y     ) & 0xff;
    secureWrite(sock, ptrEv, 6);
}

//  secureRead: 読み残しのない read
int VNC::secureRead (int s, unsigned char *data, int len) {
    //  ソケット読み出し
    int res = read(s, data, len);
    //  0: EOF, -1: エラー
    if (res <= 0) {
        printf("VNC (error): cannot read from socket (VNC)\n");
        exit(1);
    }
    //  読み残しがなければ終了，あれば再帰！
    if (res == len)
        return len;
    else
        return res + secureRead(s, data + res, len - res);
}
//  secureWrite: 書き残しのない write
int VNC::secureWrite (int s, unsigned char *data, int len) {
    //  ソケット書き込み
    int res = write(s, data, len);
    //  0/-1: エラー
    if (res <= 0) {
        printf("VNC (error): cannot write to socket (VNC)\n");
        exit(1);
    }
    //  書き残しがなければ終了，あれば再帰！
    if (res == len)
        return len;
    else
        return res + secureWrite(s, data + res, len - res);
}

//