RFBProtocol.scala

/**
Copyright (c) 2009 ymnk, JCraft,Inc. All rights reserved.
    
Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:

   1. Redistributions of source code must retain the above copyright notice,
      this list of conditions and the following disclaimer.

   2. Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in
      the documentation and/or other materials provided with the distribution.

   3. The names of the authors may not be used to endorse or promote products
      derived from this software without specific prior written permission.

 THIS SOFTWARE IS PROVIDED ``AS IS'' AND ANY EXPRESSED OR IMPLIED WARRANTIES,
 INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
 FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL JCRAFT,
 INC. OR ANY CONTRIBUTORS TO THIS SOFTWARE BE LIABLE FOR ANY DIRECT, INDIRECT,
 INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
 LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA,
 OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE,
 EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

import java.net.{InetAddress, Socket, InetSocketAddress}
import java.io.{DataInputStream, DataOutputStream}

class RFBProtocol(val host:String, val port:Int){
  import RFBProtocol._

  val TIMEOUT = 3*1000

  var remoteVersion:Int = _
  var password = ""

  var shareConnection = true

  class FrameInfo(val width:Int, val height:Int,
                  val bitsPerPixel:Int,
                  val depth:Int,
                  val bitEndianFlag:Int, val trueColorFlag:Int)

  var frameInfo:FrameInfo = _

  var favoriteAuth:List[Authentication] =
    List(noneAuthentication, vncAuthentication)

  import EncodeType._
  var favoriteEncodings:List[Int] =
    List(Raw, CopyRect, Cursor, PointerPos)

  var out:DOS = _
  var in:DIS = _

  def readChar = in.readChar&0xffff
  def readByte = in.readByte&0xff
  def readBytes(n:Int) = in.readBytes(n)
  def writeByte(arg:Int*) = arg.foreach(out.writeByte(_))
  def writeChar(arg:Int*) = arg.foreach(out.writeChar(_))

  def connect = try{

    openSocket(InetAddress.getByName(host), port) match{
      case(i, o) => in=i; out=o
    }

    remoteVersion = exchangeProtocolVersion

    securityCheck

    clientInit
    frameInfo = serverInit

    setEncodings

    frameBufferUpdateRequest(false,
                             0, 0,
                             frameInfo.width, frameInfo.height)

    while(true){
      import MessageTypeS2C._
      var command = readByte

      command match{

        case FramebufferUpdate =>{
          in.skipBytes(1)  // padding
          var numberOfRectangles = readChar
          while(numberOfRectangles > 0){
            val x = readChar
            val y = readChar
            val width =  readChar
            val height = readChar
            val encodingType = in.readInt


            import EncodeType._
            encodingType match{
              case Raw =>{
                in.skipBytes(width*height*frameInfo.bitsPerPixel/8)
                println("update: Raw %d %d %d %d".format(x, y, width, height))
              }
              case CopyRect =>{
                val srcX = readChar
                val srcY = readChar
                println("update: CopyRect %d %d %d %d %d %d".
                        format(srcX, srcY, x, y, width, height))
              }
              case Cursor =>{
                in.skipBytes(width*height*frameInfo.bitsPerPixel/8)
                in.skipBytes(height*Math.floor((width+7)/8).asInstanceOf[Int])
                println("update: Cursor")
              }
              case PointerPos =>{
                println("update: PointerPos %d %d".format(x, y))
              }
              case  _ =>
                println("unknown encodingType %d".format(encodingType))
            }
            numberOfRectangles -= 1
          }

          frameBufferUpdateRequest(true,
                                   0, 0, frameInfo.width, frameInfo.height)
        }

        case SetColourMapEntries =>{
          in.skipBytes(1) // padding
          val firstColour = readChar
          var numberOfColours = readChar
          while(numberOfColours > 0){
            val red = readChar
            val green = readChar
            val blue = readChar
            numberOfColours -= 1
          }
        }

        case Bell =>{
          System.out.println("bell")
        }

        case ServerCutText =>{
          in.skipBytes(3) // padding
          val length = in.readInt
          val text = readBytes(length)
          System.out.println("ServerCutText: "+new String(text))
        }

        case _ =>
      }
    }
  }
  catch{
    case e => println(e)
  }

  def exchangeProtocolVersion = {
    val versionString = RFBProtocol.versionString
    val protocolVersion = readBytes(12)
    val remoteVersion = new String(protocolVersion) match{
      case versionString(n) => 30+n.toInt
      case _ => 0
    }
    out.write(protocolVersion); out.flush
    remoteVersion
  }

  def securityCheck(){
    val number_of_security_types = readByte
    val security_types=new Array[Byte](number_of_security_types)
    in.readFully(security_types)

    val ERROR = -1
    val request_security_type =
      (ERROR /: favoriteAuth.map(_.typ)){
        case(b@ERROR, n) => if(security_types.exists(_==n)) n else b
        case (b, _) => b
      }

    if(request_security_type != ERROR){
      out.writeByte(request_security_type); out.flush
      val List(auth) = favoriteAuth.filter(_.typ==request_security_type)
      auth(this)
    }
    else{
      println("failed to negociate about auth type.")
    }
  }

  def clientInit{
    out.writeByte(if(shareConnection) 1 else 0 )
    out.flush
  }

  def serverInit = {
    val width = readChar
    val height = readChar
    val bitsPerPixel = readByte
    val depth = readByte
    val bitEndianFlag = readByte
    val trueColorFlag = readByte

    val (redMask, greenMask, blueMask) = (readChar, readChar, readChar)
    val (redShift, greenShift, blueShift) = (readByte, readByte, readByte)
    in.skipBytes(3)
    val nameLength = in.readInt
    val name = readBytes(nameLength)

    new FrameInfo(width, height, bitsPerPixel, depth,
                  bitEndianFlag, trueColorFlag)
  }

  def setEncodings{
    writeByte(MessageTypeC2S.SetEncodings, 0)
    out.writeChar(favoriteEncodings.size)
    favoriteEncodings.foreach(out.writeInt(_))
    out.flush
  }

  def frameBufferUpdateRequest(incremental:Boolean,
                               x:Int, y:Int,
                               width:Int, height:Int){
    writeByte(MessageTypeC2S.FramebufferUpdateRequest,
              if(incremental) 1 else 0)
    writeChar(x, y, width, height)
    out.flush
  }


  private def openSocket(ip:InetAddress, port:Int) = {
    val s = new Socket
    s.connect(new InetSocketAddress(ip, port), TIMEOUT)
    (new DataInputStream(s.getInputStream) with ReadBytes,
     new DataOutputStream(s.getOutputStream))
  }
}

object RFBProtocol{
  val versionString="RFB 003.00([1-9])\\n".r

  trait ReadBytes{ self:DataInputStream  =>
    def readBytes(n:Int) = {
      val buf=new Array[Byte](n)
      self.readFully(buf)
      buf
    }
  }

  type DOS = DataOutputStream
  type DIS = DataInputStream with ReadBytes

  object MessageTypeC2S{
    val SetPixelFormat = 0
    val SetEncodings = 2
    val FramebufferUpdateRequest = 3
    val KeyEvent = 4
    val PointerEvent = 5
    val ClientCutText = 6
  }


  object MessageTypeS2C{
    val FramebufferUpdate = 0
    val SetColourMapEntries = 1
    val Bell = 2
    val ServerCutText = 3
  }

  object SecurityType{
    val Invalid = 0
    val None = 1
    val VNCAuthentication = 2
  }

  object EncodeType{
    val Raw = 0
    val CopyRect = 1
    val RRE = 2
    val Hextile = 5
    val ZRLE = 16
    val Cursor = -239
    val PointerPos = -232
  }

  trait Authentication{
    def typ:Int
    def apply(rfb:RFBProtocol):Int
    def securityResult(rfb:RFBProtocol):Int = rfb.in.readInt
  }

  object noneAuthentication extends Authentication{
    def typ = SecurityType.None
    def apply(rfb:RFBProtocol):Int =
      if(rfb.remoteVersion>=38){ securityResult(rfb) }else{ 0 }
  }

  object vncAuthentication extends Authentication{
    import javax.crypto._
    import javax.crypto.spec._
    import java.security._

    // generate an array: 0x01 0x02 0x04 0x08 0x10 0x20 0x40 0x80
    private lazy val bits = for(i<-(0 until 8)) yield 1<<i

    def typ = SecurityType.VNCAuthentication

    def apply(rfb:RFBProtocol):Int = {
      import rfb.{in, out, password}

      val challenge = in.readBytes(16)
      val key = new Array[Byte](8) match{
        case key =>
           val p = password.getBytes
           val plen = if(p.length<key.length) p.length else key.length
           System.arraycopy(p, 0, key, 0, plen); key
      }


      // Invert every byte.  For example,
      //   0x20 (0010 0000) should be converted to 0x04 (0000 0100) .
      val _key=for(x<-key)
                 yield
                 (0 /: bits){(b, i) => b<<1|(if((x&i)==i) 1 else 0)}
                 .asInstanceOf[Byte]

      val keyFac = SecretKeyFactory.getInstance("DES")
      val secKey = keyFac.generateSecret(new DESKeySpec(_key))
      val c = Cipher.getInstance("DES/ECB/NoPadding")
      c.init(Cipher.ENCRYPT_MODE, secKey);
      out.write(c.doFinal(challenge)); out.flush
      securityResult(rfb)
    }
  }
}

object RFBProtocolMain{
  def main(arg:Array[String]){
    val Array(host, port, password@_*) = arg
    val rfbp=new RFBProtocol(host, port.toInt)
    if(password.length>0){
      rfbp.password = password.first
    }
    rfbp.connect
  }
}