vncrec 1.0.1

data/lib/vncrec/rfb/enczrle.rb

@@ -0,0 +1,202 @@
+module VNCRec
+  module RFB
+
+    module EncZRLE
+
+      class Stream
+
+        def initialize(io,bitspp,depth)
+          @io = io
+          @zstream = Zlib::Inflate.new
+          @bpp_orig = (bitspp.to_f/8.0).ceil
+          @bpp = case bitspp
+                 when 32 then 
+                   @depth <= 24 ? 3 : 4
+                 when 8 then
+                   1
+                 else
+                   raise "Cannot handle such pixel format"
+                 end
+          @depth = depth
+        end
+
+        def read_zchunk
+          zdata_len = (@io.readpartial 4).unpack("L>")[0]
+          zdata = ""
+          to_read = zdata_len
+
+          while zdata.size < zdata_len
+            zdata += @io.read(to_read)
+            to_read = zdata_len - zdata.size
+          end
+          return zdata
+        end
+
+        def read_rect(w,h)
+          fb = Array.new(w*h*@bpp)
+          data = ""
+          begin
+            data = @zstream.inflate(read_zchunk)
+          rescue Zlib::DataError
+            return
+          end
+
+          stream = StringIO.new data
+
+          tile_cols = (w.to_f/64).ceil
+          tile_rows = (h.to_f/64).ceil
+          tile_cols_rem = w % 64
+          tile_rows_rem = h % 64
+
+          tile_rows.times do |tile_row_num|
+            tile_cols.times do |tile_col_num|
+              th = if ((tile_row_num == tile_rows-1) and (tile_rows_rem > 0)) 
+                     tile_rows_rem
+                   else
+                     64
+                   end
+              tw = if ((tile_col_num == tile_cols-1) and (tile_cols_rem > 0))
+                     tile_cols_rem 
+                   else
+                     64
+                   end
+
+              subenc = stream.readbyte
+              tile = case subenc
+                     when 0 then 		#Raw
+                       tile = Array.new(tw*th)
+                       th.times do
+                         tile << (stream.read tw*@bpp_orig).unpack("C*").join
+                       end
+                       tile
+                     when 1 then 		#Solid
+                       Array.new(tw*th, stream.read(@bpp_orig))
+                     when 2..16 then 		#Packed palette
+                       handle_ZRLE_packed_palette(stream, subenc, tw,th)
+                     when 128 then 		#Plain RLE
+                       handle_ZRLE_plain_RLE_tile(stream,tw,th)
+                     when 130..255 then 	#RLE w/ palette
+                       handle_ZRLE_palette_RLE_tile(stream,subenc-128, tw,th)
+                     end
+
+              th.times do |y|
+                boline = (64 * tile_row_num + y) * w
+                offx = 64 * tile_col_num
+                fb[(boline+offx)...(boline+offx+tw)] = tile[(tw*y)...(tw*(y+1))]
+              end
+            end  #tile_col.times
+          end#tile_row.times
+          return fb.join
+        end
+
+        def handle_ZRLE_palette_RLE_tile(stream,psize,tw=64,th=64)
+          palette = []
+          pixels = Array.new(tw*th)
+          psize.times do
+            palette << stream.read(@bpp)
+          end
+          len = 0
+          begin
+            while len < tw*th
+              id = stream.read(1).unpack("C")[0]
+              #
+              #+--------+--------+--------+--------+
+              #|   id   |   255  |   ..   |  <255  |
+              #+--------+--------+--------+--------+
+              #
+              if (id & 0b10000000) == 0 
+                rl = 1
+              else
+                id -= 128
+                rl = 0
+                rem = 0
+                while (rem = stream.readbyte) == 255
+                  rl += 255
+                end
+                rl += rem + 1 
+              end
+              pixels[len...(len+rl)] = Array.new(rl, palette[id]) #TODO: if rl == 1 
+              len += rl
+            end
+          rescue EOFError
+          end
+          pixels
+        end
+
+        def handle_ZRLE_plain_RLE_tile(stream,tw=64,th=64)
+          pixels = Array.new(tw*th)
+          len = 0
+          begin
+            while len < tw*th
+              color = stream.read(@bpp)
+              #
+              #+--------+--------+--------+--------+
+              #|  color |   255  |   ..   |  <255  |
+              #+--------+--------+--------+--------+
+              #
+              rl = 0
+              rem = 0
+              while (rem = stream.readbyte) == 255
+                rl += 255
+              end
+              rl += rem + 1 
+              pixels[len...(len+rl)] = Array.new(rl, color) #TODO: if rl == 1 
+              len += rl
+            end
+          rescue EOFError
+          end
+          pixels
+        end
+
+        def handle_ZRLE_packed_palette(stream, psize, tw=64, th=64)
+          pixels = Array.new(tw*th, 0)
+          bitspp = case psize
+                   when 2 then 1
+                   when 3..4 then 2
+                   when 5..16 then 4
+                   else 
+                     return pixels
+                   end
+          palette = []
+          psize.times do
+            palette << stream.read(@bpp).unpack("C*")
+          end
+          count = case psize
+                  when 2 then th*((tw+7)/8)
+                  when 3..4 then th*((tw+3)/4)
+                  when 5..16 then th*((tw+1)/2)
+                  end
+          off_bits = 0
+          bits_per_row = bitspp * tw
+          padding_bits = bits_per_row % 8
+          encoded_len_bits = 64 * (bits_per_row + padding_bits)
+          encoded = stream.read(count).unpack("C*")
+          pixnum = 0
+          while off_bits < (encoded_len_bits - padding_bits - bitspp)
+            b1 = encoded[off_bits/8]
+            b2 = encoded[off_bits/8 + 1] || 0
+            b1 <<= 8
+            pixels[pixnum] = palette[h_bitmask(b2 + b1, bitspp, off_bits % 16)].pack("C*")
+            off_bits += if (off_bits % bits_per_row) > (bits_per_row - padding_bits) and (padding_bits > 0)
+                          bitspp + padding_bits
+                        else
+                          bitspp
+                        end
+            pixnum += 1
+          end
+          pixels
+        end
+      end
+
+
+    end
+
+  end
+end
+
+def h_bitmask(input,count,offset=0)
+  #return first n bits of ushort as integer
+  #TODO: make generalization of input type
+  input <<= offset
+  (input & (0xFFFF - 2**(16-count) + 1)) >> (16 - count)
+end

data/lib/vncrec/rfb/proxy.rb

@@ -0,0 +1,200 @@
+require 'socket'
+
+require 'vncrec/rfb/encraw.rb'
+require 'vncrec/rfb/enczrle.rb'
+require 'vncrec/rfb/enchex.rb'
+
+module VNCRec
+  module RFB
+    class Proxy
+      attr_accessor :name, :w, :h, :io, :data
+
+      # @param io [IO, #read, #sysread, #syswrite, #read_nonblock] string stream from VNC server.
+      # @param rfbv [String] version of RFB protocol, 3.8 is the only supported by now
+      # @param enc [Integer] encoding of video data used to transfer. One of the following:
+      #  * {ENC_RAW}
+      #  * {ENC_HEXTILE}
+      #  * {ENC_ZRLE}
+      # @param pf [Hash] pixel format:
+      #  * {VNCRec::PIX_FMT_BGR8} - 8 bits per pixel
+      #  * {VNCRec::PIX_FMT_BGR32} - 32 bits per pixel
+      # @param w width of the screen area
+      # @param h height of the screen area
+      def initialize(io, rfbv, enc, pf)
+        @io = io
+        @version = rfbv
+        @enc = enc
+        @pf = pf
+      end
+
+      def prepare_framebuffer(w, h, bpp)
+        @w = w
+        @h = h
+        @bpp = bpp
+        @bypp = (bpp / 8.0).to_i
+        @wb = @w * @bypp
+        @data = "\x00" * @wb * @h
+      end
+
+      # Perform handshake
+      # @return w,h,server_name or nil
+      def handshake
+        # version
+        version = @io.readpartial 12
+        @io.syswrite(@version + "\n")
+
+        # security
+        num_of_st = @io.readbyte
+        if num_of_st == 0 # failed
+          reason_len = @io.readpartial(4).unpack('L>')[0]
+          reason = @io.readpartial(reason_len)
+          fail reason
+        else
+          num_of_st.times do
+            @io.readbyte
+          end
+        end
+
+        reply = [1].pack('C') # security type:none
+        @io.syswrite reply
+
+        stype = (@io.readpartial 4).unpack('H' * 8)
+        # client init
+        @io.syswrite reply
+
+        # server init
+        w = @io.readpartial(2).unpack('S>')[0]
+        h = @io.readpartial(2).unpack('S>')[0]
+        pf = @io.readpartial 16
+        nlen = @io.readpartial(4).unpack('L>')[0]
+        @name = @io.readpartial nlen
+        return [w, h, @name]
+      rescue
+        return nil
+      end
+
+      # Set a way that server should use to represent pixel data
+      # @param [Hash] pixel format:
+      #  * {VNCRec::PIX_FMT_BGR8}
+      #  * {VNCRec::PIX_FMT_BGRA}
+      def set_pixel_format(format)
+        msg = [0, 0, 0, 0].pack('CC3')
+        begin
+          @io.syswrite msg
+
+          msg = [
+            format[:bpp],
+            format[:depth],
+            format[:bend],
+            format[:tcol],
+            format[:rmax],
+            format[:gmax],
+            format[:bmax],
+            format[:rshif],
+            format[:gshif],
+            format[:bshif],
+            0, 0, 0
+          ].pack('CCCCS>S>S>CCCC3')
+          return @io.syswrite msg
+
+        rescue
+          return nil
+        end
+      end
+
+      # Set way of encoding video frames.
+      # @param encodings [Array<Integer>] encoding of video data used to transfer.
+      #  * {ENC_RAW}
+      #  * {ENC_HEXTILE}
+      #  * {ENC_ZRLE}
+      def set_encodings(encodings)
+        num = encodings.size
+        msg = [2, 0, num].pack('CCS>')
+        begin
+          @io.syswrite msg
+          encodings.each do |e|
+            @io.syswrite([e].pack('l>'))
+          end
+        rescue
+          return nil
+        end
+      end
+
+      # Request framebuffer update.
+      # @param [Integer] inc incremental, request just difference
+      #  between previous and current framebuffer state.
+      # @param x [Integer]
+      # @param y [Integer]
+      # @param w [Integer]
+      # @param h [Integer]
+      def fb_update_request(inc, x, y, w, h)
+        @inc = inc > 0
+        msg = [3, inc, x, y, w, h].pack('CCS>S>S>S>')
+        return @io.write msg
+      rescue
+        return nil
+      end
+
+      # Handle VNC server response. Call it right after +fb_update_request+.
+      # @return [Array] type, (either framebuffer, "bell", +handle_server_cuttext+ or +handle_colormap_update+ results)
+      def handle_response
+        t = (io.readpartial 1).ord
+        case t
+        when 0 then
+          handle_fb_update
+          return [t, @data]
+        when 1 then
+          return [t, handle_colormap_update]
+        when 2 then
+          return [t, 'bell']
+        when 3 then
+          return [t, handle_server_cuttext]
+        else
+          return [-1, nil]
+        end
+      end
+
+      # Receives data and applies diffs(if incremental) to the @data
+      def handle_fb_update
+        fail 'run #prepare_framebuffer first' unless @data
+        enc = nil
+        @encs ||= { 0 => VNCRec::RFB::EncRaw,
+                    5 => VNCRec::RFB::EncHextile,
+                    16 => VNCRec::RFB::EncZRLE
+        }
+        _, numofrect = @io.read(3).unpack('CS>')
+        i = 0
+        while i < numofrect
+          hdr = @io.read 12
+          x, y, w, h, enc = hdr.unpack('S>S>S>S>l>')
+          mod = @encs.fetch(enc) { fail "Unsupported encoding #{enc}" }
+          mod.read_rect @io, x, y, w, h, @bpp, @data, @wb, @h
+          i += 1
+        end
+      end
+
+      # @return [Array] palette
+      def handle_colormap_update
+        _, first_color, noc = (@io.read 5).unpack('CS>S>')
+        palette = []
+        noc.times do
+          palette << (@io.read 6).unpack('S>S>S>')
+        end
+        return palette
+      rescue
+        return nil
+      end
+
+      # @return [String] server cut text
+      def handle_server_cuttext
+        begin
+          _, _, _, len = (@io.read 7).unpack('C3L>')
+          text = @io.read len
+        rescue
+          return nil
+        end
+        text
+      end
+    end
+  end
+end

data/lib/vncrec/version.rb

@@ -0,0 +1,3 @@
+module VNCRec
+  VERSION = "1.0.1"
+end

data/lib/vncrec/writers.rb

@@ -0,0 +1,209 @@
+require "vncrec/constants.rb"
+require "timeout"
+require "thread"
+
+module VNCRec
+  # Writers are wrappers for video files.
+  module Writers
+    # Raw video writer. Very similar to File
+    class RawVideo
+      def initialize(filename)
+        @filename = filename
+        @file = File.open(filename, 'w')
+      end
+
+      def write(data)
+        @file.write data
+        @file.flush
+      end
+
+      def close
+        @file.close
+      end
+
+      def closed?
+        @file.closed?
+      end
+
+      def size
+        @file.size
+      end
+    end
+
+    # FFmpeg writer. Pipes video to FFmpeg instance exactly
+    # *fps* times per second. Audio addition is also
+    # supported (`:ffmpeg_ia`- and `:ffmpeg_out_opts strings`)
+    class FFmpeg
+      # @param filename [String] a name for video file.
+      #  Should contain extension i.e. _.mp4_ of _.flv_.
+      # @note Choose _-acodec_ option in +:ffmpeg_out_opts+ accordingly.
+      # @param opts [Hash] options:
+      #  * fps
+      #  * pix_fmt (see +:colormode+)
+      #  * geometry
+      #  * ffmpeg_iv_opts
+      #  * ffmpeg_ia_opts
+      #  * ffmpeg_out_opts
+      #  See {VNCRec::Recorder#initialize} for descriptions
+      def initialize(filename, opts = {})
+        @filename = filename
+        @fps = opts[:fps] || 12
+        pf = opts.fetch(:pix_fmt) { fail 'Undefined pixel format' }
+        @pix_fmt = get_pix_fmt pf
+        @size = opts.fetch(:geometry) { fail 'Undefined frame size' }
+        @frame_length = frame_length
+        @ffmpeg_iv_opts = opts[:ffmpeg_iv_opts]
+        @ffmpeg_ia_opts = opts[:ffmpeg_ia_opts]
+        @ffmpeg_out_opts = opts[:ffmpeg_out_opts]
+        @cmd = "ffmpeg -y -s #{@size} -r #{@fps} -f rawvideo -pix_fmt #{@pix_fmt[:string]} \
+        #{@ffmpeg_iv_opts} \
+      -i pipe:0 \
+        #{@ffmpeg_ia_opts} \
+        #{@ffmpeg_out_opts} #{@filename} &>/dev/null"
+        @data_avail = false
+        spawn
+      end
+
+      def write(data)
+        begin
+          written = @pipe_to_writer.syswrite(data)
+        rescue Errno::EPIPE
+          raise 'No writer running'
+        end
+        fail 'Not enough data is piped to writer' if written % @frame_length != 0
+        @pipe_to_writer.flush
+        @data_avail = true
+      end
+
+      def close
+        Process.kill('KILL', @pid)
+        Timeout.timeout(5) do
+          Process.waitpid(@pid)
+        end
+      rescue Timeout::Error
+        raise 'Writer hanged'
+      rescue Errno::ESRCH, Errno::ECHILD
+        raise 'No writer running'
+      end
+
+      def closed?
+        Timeout.timeout(0.05) do
+          Process.waitpid(@pid)
+          return true
+        end
+      rescue Timeout::Error
+        return false
+      rescue Errno::ECHILD
+        return true
+      end
+
+      # @return [Integer] filesize. If no file created yet
+      #  0 is returned.
+      def size
+        s = File.size(@filename)
+        return s
+      rescue Errno::ENOENT
+        return 0
+      end
+
+      private
+
+      def spawn
+        @pipe, @pipe_to_writer = IO.pipe
+        @pid = fork do
+          Signal.trap('INT') {}
+          @pipe_to_writer.close
+          @lock = Mutex.new
+          @written = 0
+          @output_ready = false
+          STDIN.reopen(@pipe)
+          routine
+        end
+        @pipe.close
+      end
+
+      def routine
+        @output = IO.popen(@cmd)
+        @output_ready = true
+        IO.select([STDIN])
+        @th = Thread.new(thread_func)
+        loop do
+          data = STDIN.read(@frame_length)
+          fail 'wrong length' if data.length != @frame_length
+          if @lock.try_lock
+            @framebuffer = data
+            @lock.unlock
+          else
+            @framebuffer2 = data
+            @cached = true
+          end
+        end
+      end
+
+      def flush
+        return unless @output.closed? || @data_avail
+        @lock.synchronize do
+          if @cached
+            @framebuffer = @framebuffer2
+            @cached = false
+          end
+          @written += @output.syswrite @framebuffer
+          @output.flush
+        end
+      end
+
+      def thread_func
+        adjust_sleep_time { flush }
+        i = 0
+        loop do
+          i += 1
+          if (i % 100) == 0
+            adjust_sleep_time { flush }
+          else
+            flush
+          end
+          sleep @sl
+        end
+      end
+
+      def frame_length
+        bpp = @pix_fmt[:bpp] / 8
+        dim = @size.split('x').map(&:to_i).reduce(&:*)
+        bpp * dim
+      end
+
+      def adjust_sleep_time(&_block)
+        t1 = Time.now
+        yield
+        t2 = Time.now
+        @sl = 1.0 / @fps - (t2 - t1)
+      end
+
+      def get_pix_fmt(fmt)
+        sym = fmt.to_s.upcase.prepend('PIX_FMT_').to_sym
+        fail "Unknown pixel format #{fmt}" unless VNCRec.const_defined? sym
+        VNCRec.const_get(sym)
+      end
+    end
+
+    def self.get_writer(filename, opts = {})
+      begin
+        File.write(filename, '')
+      rescue Errno::EACCES
+        raise 'Cannot create output file'
+      end
+      @path, @filename = File.split filename
+      @extname = File.extname filename
+      return RawVideo.new(@path + '/' + @filename) if @extname == '.raw'
+      if @extname.empty?
+        if @path != '/dev'
+          return RawVideo.new(@path + '/' + @filename + '.raw')
+        else
+          return FFmpeg.new(@path + '/' + @filename, opts)
+        end
+      else
+        return FFmpeg.new(@path + '/' + @filename, opts)
+      end
+    end
+  end
+end