diff --git a/bin/vnc_recorder b/bin/vnc_recorder new file mode 100755 index 0000000..c54e0e5 Binary files /dev/null and b/bin/vnc_recorder differ diff --git a/vnc_rec/message-listeners.go b/vnc_rec/message-listeners.go new file mode 100644 index 0000000..7acc950 --- /dev/null +++ b/vnc_rec/message-listeners.go @@ -0,0 +1,77 @@ +package vnc_rec + +import ( + "github.com/amitbet/vncproxy/client" + "github.com/amitbet/vncproxy/common" + "github.com/amitbet/vncproxy/logger" + "github.com/amitbet/vncproxy/server" +) + +type ClientUpdater struct { + conn *client.ClientConn +} + +// Consume recieves vnc-server-bound messages (Client messages) and updates the server part of the proxy +func (cc *ClientUpdater) Consume(seg *common.RfbSegment) error { + logger.Tracef("ClientUpdater.Consume (vnc-server-bound): got segment type=%s bytes: %v", seg.SegmentType, seg.Bytes) + switch seg.SegmentType { + + case common.SegmentFullyParsedClientMessage: + clientMsg := seg.Message.(common.ClientMessage) + logger.Debugf("ClientUpdater.Consume:(vnc-server-bound) got ClientMessage type=%s", clientMsg.Type()) + switch clientMsg.Type() { + + case common.SetPixelFormatMsgType: + // update pixel format + logger.Debugf("ClientUpdater.Consume: updating pixel format") + pixFmtMsg := clientMsg.(*server.MsgSetPixelFormat) + cc.conn.PixelFormat = pixFmtMsg.PF + } + + err := clientMsg.Write(cc.conn) + if err != nil { + logger.Errorf("ClientUpdater.Consume (vnc-server-bound, SegmentFullyParsedClientMessage): problem writing to port: %s", err) + } + return err + } + return nil +} + +type ServerUpdater struct { + conn *server.ServerConn +} + +func (p *ServerUpdater) Consume(seg *common.RfbSegment) error { + + logger.Debugf("WriteTo.Consume (ServerUpdater): got segment type=%s, object type:%d", seg.SegmentType, seg.UpcomingObjectType) + switch seg.SegmentType { + case common.SegmentMessageStart: + case common.SegmentRectSeparator: + case common.SegmentServerInitMessage: + serverInitMessage := seg.Message.(*common.ServerInit) + p.conn.SetHeight(serverInitMessage.FBHeight) + p.conn.SetWidth(serverInitMessage.FBWidth) + p.conn.SetDesktopName(string(serverInitMessage.NameText)) + p.conn.SetPixelFormat(&serverInitMessage.PixelFormat) + + case common.SegmentBytes: + logger.Debugf("WriteTo.Consume (ServerUpdater SegmentBytes): got bytes len=%d", len(seg.Bytes)) + _, err := p.conn.Write(seg.Bytes) + if err != nil { + logger.Errorf("WriteTo.Consume (ServerUpdater SegmentBytes): problem writing to port: %s", err) + } + return err + case common.SegmentFullyParsedClientMessage: + + clientMsg := seg.Message.(common.ClientMessage) + logger.Debugf("WriteTo.Consume (ServerUpdater): got ClientMessage type=%s", clientMsg.Type()) + err := clientMsg.Write(p.conn) + if err != nil { + logger.Errorf("WriteTo.Consume (ServerUpdater SegmentFullyParsedClientMessage): problem writing to port: %s", err) + } + return err + default: + //return errors.New("WriteTo.Consume: undefined RfbSegment type") + } + return nil +} diff --git a/vnc_rec/proxy.go b/vnc_rec/proxy.go new file mode 100644 index 0000000..8185552 --- /dev/null +++ b/vnc_rec/proxy.go @@ -0,0 +1,203 @@ +package vnc_rec + +import ( + "net" + "path" + "strconv" + "time" + + "github.com/amitbet/vncproxy/client" + "github.com/amitbet/vncproxy/common" + "github.com/amitbet/vncproxy/encodings" + "github.com/amitbet/vncproxy/logger" + "github.com/amitbet/vncproxy/player" + "github.com/amitbet/vncproxy/server" +) + +type VncProxy struct { + TCPListeningURL string // empty = not listening on tcp + WsListeningURL string // empty = not listening on ws + RecordingDir string // empty = no recording + ProxyVncPassword string //empty = no auth + SingleSession *VncSession // to be used when not using sessions + UsingSessions bool //false = single session - defined in the var above + sessionManager *SessionManager +} + +func (vp *VncProxy) createClientConnection(target string, vncPass string) (*client.ClientConn, error) { + var ( + nc net.Conn + err error + ) + + if target[0] == '/' { + nc, err = net.Dial("unix", target) + } else { + nc, err = net.Dial("tcp", target) + } + + if err != nil { + logger.Errorf("error connecting to vnc server: %s", err) + return nil, err + } + + var noauth client.ClientAuthNone + authArr := []client.ClientAuth{&client.PasswordAuth{Password: vncPass}, &noauth} + + clientConn, err := client.NewClientConn(nc, + &client.ClientConfig{ + Auth: authArr, + Exclusive: true, + }) + + if err != nil { + logger.Errorf("error creating client: %s", err) + return nil, err + } + + return clientConn, nil +} + +// if sessions not enabled, will always return the configured target server (only one) +func (vp *VncProxy) getProxySession(sessionId string) (*VncSession, error) { + + if !vp.UsingSessions { + if vp.SingleSession == nil { + logger.Errorf("SingleSession is empty, use sessions or populate the SingleSession member of the VncProxy struct.") + } + return vp.SingleSession, nil + } + return vp.sessionManager.GetSession(sessionId) +} + +func (vp *VncProxy) newServerConnHandler(cfg *server.ServerConfig, sconn *server.ServerConn) error { + var err error + session, err := vp.getProxySession(sconn.SessionId) + if err != nil { + logger.Errorf("Proxy.newServerConnHandler can't get session: %d", sconn.SessionId) + return err + } + + var rec *Recorder + + if session.Type == SessionTypeRecordingProxy { + recFile := "recording" + strconv.FormatInt(time.Now().Unix(), 10) + ".rbs" + recPath := path.Join(vp.RecordingDir, recFile) + rec, err = NewRecorder(recPath) + if err != nil { + logger.Errorf("Proxy.newServerConnHandler can't open recorder save path: %s", recPath) + return err + } + + sconn.Listeners.AddListener(rec) + } + + session.Status = SessionStatusInit + if session.Type == SessionTypeProxyPass || session.Type == SessionTypeRecordingProxy { + target := session.Target + if session.TargetHostname != "" && session.TargetPort != "" { + target = session.TargetHostname + ":" + session.TargetPort + } + + cconn, err := vp.createClientConnection(target, session.TargetPassword) + if err != nil { + session.Status = SessionStatusError + logger.Errorf("Proxy.newServerConnHandler error creating connection: %s", err) + return err + } + if session.Type == SessionTypeRecordingProxy { + cconn.Listeners.AddListener(rec) + } + + //creating cross-listeners between server and client parts to pass messages through the proxy: + + // gets the bytes from the actual vnc server on the env (client part of the proxy) + // and writes them through the server socket to the vnc-client + serverUpdater := &ServerUpdater{sconn} + cconn.Listeners.AddListener(serverUpdater) + + // gets the messages from the server part (from vnc-client), + // and write through the client to the actual vnc-server + clientUpdater := &ClientUpdater{cconn} + sconn.Listeners.AddListener(clientUpdater) + + err = cconn.Connect() + if err != nil { + session.Status = SessionStatusError + logger.Errorf("Proxy.newServerConnHandler error connecting to client: %s", err) + return err + } + + encs := []common.IEncoding{ + &encodings.RawEncoding{}, + &encodings.TightEncoding{}, + &encodings.EncCursorPseudo{}, + &encodings.EncLedStatePseudo{}, + &encodings.TightPngEncoding{}, + &encodings.RREEncoding{}, + &encodings.ZLibEncoding{}, + &encodings.ZRLEEncoding{}, + &encodings.CopyRectEncoding{}, + &encodings.CoRREEncoding{}, + &encodings.HextileEncoding{}, + } + cconn.Encs = encs + + if err != nil { + session.Status = SessionStatusError + logger.Errorf("Proxy.newServerConnHandler error connecting to client: %s", err) + return err + } + } + + if session.Type == SessionTypeReplayServer { + fbs, err := player.ConnectFbsFile(session.ReplayFilePath, sconn) + + if err != nil { + logger.Error("TestServer.NewConnHandler: Error in loading FBS: ", err) + return err + } + sconn.Listeners.AddListener(player.NewFBSPlayListener(sconn, fbs)) + return nil + + } + + session.Status = SessionStatusActive + return nil +} + +func (vp *VncProxy) StartListening() { + + secHandlers := []server.SecurityHandler{&server.ServerAuthNone{}} + + if vp.ProxyVncPassword != "" { + secHandlers = []server.SecurityHandler{&server.ServerAuthVNC{vp.ProxyVncPassword}} + } + cfg := &server.ServerConfig{ + SecurityHandlers: secHandlers, + Encodings: []common.IEncoding{&encodings.RawEncoding{}, &encodings.TightEncoding{}, &encodings.CopyRectEncoding{}}, + PixelFormat: common.NewPixelFormat(32), + ClientMessages: server.DefaultClientMessages, + DesktopName: []byte("workDesk"), + Height: uint16(768), + Width: uint16(1024), + NewConnHandler: vp.newServerConnHandler, + UseDummySession: !vp.UsingSessions, + } + + if vp.TCPListeningURL != "" && vp.WsListeningURL != "" { + logger.Infof("running two listeners: tcp port: %s, ws url: %s", vp.TCPListeningURL, vp.WsListeningURL) + + go server.WsServe(vp.WsListeningURL, cfg) + server.TcpServe(vp.TCPListeningURL, cfg) + } + + if vp.WsListeningURL != "" { + logger.Infof("running ws listener url: %s", vp.WsListeningURL) + server.WsServe(vp.WsListeningURL, cfg) + } + if vp.TCPListeningURL != "" { + logger.Infof("running tcp listener on port: %s", vp.TCPListeningURL) + server.TcpServe(vp.TCPListeningURL, cfg) + } +} diff --git a/vnc_rec/proxy_test.go b/vnc_rec/proxy_test.go new file mode 100644 index 0000000..aa8723b --- /dev/null +++ b/vnc_rec/proxy_test.go @@ -0,0 +1,27 @@ +package vnc_rec + +import "testing" + +func TestProxy(t *testing.T) { + //create default session if required + t.Skip("this isn't an automated test, just an entrypoint for debugging") + + proxy := &VncProxy{ + WsListeningURL: "http://0.0.0.0:7778/", // empty = not listening on ws + RecordingDir: "d:\\", // empty = no recording + TCPListeningURL: ":5904", + //RecordingDir: "C:\\vncRec", // empty = no recording + ProxyVncPassword: "1234", //empty = no auth + SingleSession: &VncSession{ + TargetHostname: "192.168.1.101", + TargetPort: "5901", + TargetPassword: "123456", + ID: "dummySession", + Status: SessionStatusInit, + Type: SessionTypeRecordingProxy, + }, // to be used when not using sessions + UsingSessions: false, //false = single session - defined in the var above + } + + proxy.StartListening() +} diff --git a/vnc_rec/recorder.go b/vnc_rec/recorder.go new file mode 100644 index 0000000..53a4f54 --- /dev/null +++ b/vnc_rec/recorder.go @@ -0,0 +1,215 @@ +package vnc_rec + +import ( + "bytes" + "encoding/binary" + "os" + "time" + + "github.com/amitbet/vncproxy/common" + "github.com/amitbet/vncproxy/logger" + "github.com/amitbet/vncproxy/server" +) + +type Recorder struct { + //common.BytesListener + RBSFileName string + writer *os.File + //logger common.Logger + startTime int + buffer bytes.Buffer + serverInitMessage *common.ServerInit + sessionStartWritten bool + segmentChan chan *common.RfbSegment + maxWriteSize int +} + +func getNowMillisec() int { + return int(time.Now().UnixNano() / int64(time.Millisecond)) +} + +func NewRecorder(saveFilePath string) (*Recorder, error) { + //delete file if it exists + if _, err := os.Stat(saveFilePath); err == nil { + os.Remove(saveFilePath) + } + + rec := Recorder{RBSFileName: saveFilePath, startTime: getNowMillisec()} + var err error + + rec.maxWriteSize = 65535 + + rec.writer, err = os.OpenFile(saveFilePath, os.O_RDWR|os.O_CREATE, 0644) + if err != nil { + logger.Errorf("unable to open file: %s, error: %v", saveFilePath, err) + return nil, err + } + + //buffer the channel so we don't halt the proxying flow for slow writes when under pressure + rec.segmentChan = make(chan *common.RfbSegment, 100) + go func() { + for { + data := <-rec.segmentChan + rec.HandleRfbSegment(data) + } + }() + + return &rec, nil +} + +const versionMsg_3_3 = "RFB 003.003\n" +const versionMsg_3_7 = "RFB 003.007\n" +const versionMsg_3_8 = "RFB 003.008\n" + +// Security types +const ( + SecTypeInvalid = 0 + SecTypeNone = 1 + SecTypeVncAuth = 2 + SecTypeTight = 16 +) + +// func (r *Recorder) writeHeader() error { +// _, err := r.writer.WriteString("FBS 001.000\n") +// return err +// // df.write("FBS 001.000\n".getBytes()); +// } + +func (r *Recorder) writeStartSession(initMsg *common.ServerInit) error { + r.sessionStartWritten = true + desktopName := string(initMsg.NameText) + framebufferWidth := initMsg.FBWidth + framebufferHeight := initMsg.FBHeight + + //write rfb header information (the only part done without the [size|data|timestamp] block wrapper) + r.writer.WriteString("FBS 001.000\n") + + //push the version message into the buffer so it will be written in the first rbs block + r.buffer.WriteString(versionMsg_3_3) + + //push sec type and fb dimensions + binary.Write(&r.buffer, binary.BigEndian, int32(SecTypeNone)) + binary.Write(&r.buffer, binary.BigEndian, int16(framebufferWidth)) + binary.Write(&r.buffer, binary.BigEndian, int16(framebufferHeight)) + + buff := bytes.Buffer{} + //binary.Write(&buff, binary.BigEndian, initMsg.FBWidth) + //binary.Write(&buff, binary.BigEndian, initMsg.FBHeight) + binary.Write(&buff, binary.BigEndian, initMsg.PixelFormat) + buff.Write([]byte{0, 0, 0}) //padding + r.buffer.Write(buff.Bytes()) + //logger.Debugf(">>>>>>buffer for initMessage:%v ", buff.Bytes()) + + //var fbsServerInitMsg = []byte{32, 24, 0, 1, 0, byte(0xFF), 0, byte(0xFF), 0, byte(0xFF), 16, 8, 0, 0, 0, 0} + //r.buffer.Write(fbsServerInitMsg) + + binary.Write(&r.buffer, binary.BigEndian, uint32(len(desktopName))) + + r.buffer.WriteString(desktopName) + //binary.Write(&r.buffer, binary.BigEndian, byte(0)) // add null termination for desktop string + + return nil +} + +func (r *Recorder) Consume(data *common.RfbSegment) error { + //using async writes so if chan buffer overflows, proxy will not be affected + select { + case r.segmentChan <- data: + // default: + // logger.Error("error: recorder queue is full") + } + + return nil +} + +func (r *Recorder) HandleRfbSegment(data *common.RfbSegment) error { + defer func() { + if r := recover(); r != nil { + logger.Error("Recovered in HandleRfbSegment: ", r) + } + }() + + switch data.SegmentType { + case common.SegmentMessageStart: + if !r.sessionStartWritten { + logger.Debugf("Recorder.HandleRfbSegment: writing start session segment: %v", r.serverInitMessage) + r.writeStartSession(r.serverInitMessage) + } + + switch common.ServerMessageType(data.UpcomingObjectType) { + case common.FramebufferUpdate: + logger.Debugf("Recorder.HandleRfbSegment: saving FramebufferUpdate segment") + //r.writeToDisk() + case common.SetColourMapEntries: + case common.Bell: + case common.ServerCutText: + default: + logger.Warn("Recorder.HandleRfbSegment: unknown message type:" + string(data.UpcomingObjectType)) + } + case common.SegmentConnectionClosed: + r.writeToDisk() + case common.SegmentRectSeparator: + logger.Debugf("Recorder.HandleRfbSegment: writing rect") + //r.writeToDisk() + case common.SegmentBytes: + logger.Debug("Recorder.HandleRfbSegment: writing bytes, len:", len(data.Bytes)) + if r.buffer.Len()+len(data.Bytes) > r.maxWriteSize-4 { + r.writeToDisk() + } + _, err := r.buffer.Write(data.Bytes) + return err + case common.SegmentServerInitMessage: + r.serverInitMessage = data.Message.(*common.ServerInit) + case common.SegmentFullyParsedClientMessage: + clientMsg := data.Message.(common.ClientMessage) + + switch clientMsg.Type() { + case common.SetPixelFormatMsgType: + clientMsg := data.Message.(*server.MsgSetPixelFormat) + logger.Debugf("Recorder.HandleRfbSegment: client message %v", *clientMsg) + r.serverInitMessage.PixelFormat = clientMsg.PF + default: + //return errors.New("unknown client message type:" + string(data.UpcomingObjectType)) + } + + default: + //return errors.New("undefined RfbSegment type") + } + return nil +} + +func (r *Recorder) writeToDisk() error { + timeSinceStart := getNowMillisec() - r.startTime + if r.buffer.Len() == 0 { + return nil + } + + //write buff length + bytesLen := r.buffer.Len() + binary.Write(r.writer, binary.BigEndian, uint32(bytesLen)) + paddedSize := (bytesLen + 3) & 0x7FFFFFFC + paddingSize := paddedSize - bytesLen + + //logger.Debugf("paddedSize=%d paddingSize=%d bytesLen=%d", paddedSize, paddingSize, bytesLen) + //write buffer padded to 32bit + _, err := r.buffer.WriteTo(r.writer) + padding := make([]byte, paddingSize) + //logger.Debugf("padding=%v ", padding) + + binary.Write(r.writer, binary.BigEndian, padding) + + //write timestamp + binary.Write(r.writer, binary.BigEndian, uint32(timeSinceStart)) + r.buffer.Reset() + return err +} + +// func (r *Recorder) WriteUInt8(data uint8) error { +// buf := make([]byte, 1) +// buf[0] = byte(data) // cast int8 to byte +// return r.Write(buf) +// } + +func (r *Recorder) Close() { + r.writer.Close() +} diff --git a/vnc_rec/session-manager.go b/vnc_rec/session-manager.go new file mode 100644 index 0000000..63c50c9 --- /dev/null +++ b/vnc_rec/session-manager.go @@ -0,0 +1,19 @@ +package vnc_rec + +type SessionManager struct { + sessions map[string]*VncSession +} + +func (s *SessionManager) GetSession(sessionId string) (*VncSession, error) { + return s.sessions[sessionId], nil +} + +func (s *SessionManager) SetSession(sessionId string, session *VncSession) error { + s.sessions[sessionId] = session + return nil +} + +func (s *SessionManager) DeleteSession(sessionId string) error { + delete(s.sessions, sessionId) + return nil +} diff --git a/vnc_rec/vnc-session.go b/vnc_rec/vnc-session.go new file mode 100644 index 0000000..9341332 --- /dev/null +++ b/vnc_rec/vnc-session.go @@ -0,0 +1,27 @@ +package vnc_rec + +type SessionStatus int +type SessionType int + +const ( + SessionStatusInit SessionStatus = iota + SessionStatusActive + SessionStatusError +) + +const ( + SessionTypeRecordingProxy SessionType = iota + SessionTypeReplayServer + SessionTypeProxyPass +) + +type VncSession struct { + Target string + TargetHostname string + TargetPort string + TargetPassword string + ID string + Status SessionStatus + Type SessionType + ReplayFilePath string +}