diff --git a/assets/index.html b/assets/index.html index 21362eb..0b0c7ba 100644 --- a/assets/index.html +++ b/assets/index.html @@ -42,9 +42,9 @@
Interactive session
-

From {{.Ip}}, running {{.Cmd}}, session ID: {{.Name}} +

From {{.Ip}}, running {{.Cmd}}, session ID: {{.Id}}

- View + View Session
diff --git a/assets/main.css b/assets/main.css index 36ee561..119fca4 100644 --- a/assets/main.css +++ b/assets/main.css @@ -35,4 +35,4 @@ .navbar-xs .navbar-nav>li>a { padding-top: 0px; padding-bottom: 0px; -} +} \ No newline at end of file diff --git a/assets/term.html b/assets/term.html index fa7268c..a39aa55 100644 --- a/assets/term.html +++ b/assets/term.html @@ -23,11 +23,11 @@ @@ -40,16 +40,33 @@ diff --git a/main.go b/main.go index db69794..a199a0e 100644 --- a/main.go +++ b/main.go @@ -6,6 +6,7 @@ import ( "net/url" "os" + "github.com/dchest/uniuri" "github.com/gin-gonic/gin" "github.com/syssecfsu/witty/term_conn" ) @@ -32,9 +33,9 @@ func checkOrigin(r *http.Request) bool { } type InteractiveSession struct { - Ip string - Cmd string - Name string + Ip string + Cmd string + Id string } func fillIndex(c *gin.Context) { @@ -42,9 +43,9 @@ func fillIndex(c *gin.Context) { term_conn.ForEachSession(func(tc *term_conn.TermConn) { players = append(players, InteractiveSession{ - Name: tc.Name, - Ip: tc.Ip, - Cmd: cmdToExec[0], + Id: tc.Name, + Ip: tc.Ip, + Cmd: cmdToExec[0], }) }) @@ -79,42 +80,60 @@ func main() { rt.SetTrustedProxies(nil) rt.LoadHTMLGlob("./assets/*.html") - rt.GET("/view/:sname", func(c *gin.Context) { - sname := c.Param("sname") - c.HTML(http.StatusOK, "term.html", gin.H{ - "title": "viewer terminal", - "path": "/ws_view/" + sname, - }) - }) - - rt.GET("/new", func(c *gin.Context) { - if host == nil { - host = &c.Request.Host - } - - c.HTML(http.StatusOK, "term.html", gin.H{ - "title": "interactive terminal", - "path": "/ws_new", - }) - }) - - rt.GET("/ws_new", func(c *gin.Context) { - term_conn.ConnectTerm(c.Writer, c.Request, false, "", cmdToExec) - }) - - rt.GET("/ws_view/:sname", func(c *gin.Context) { - path := c.Param("sname") - term_conn.ConnectTerm(c.Writer, c.Request, true, path, nil) - }) - - // handle static files - rt.Static("/assets", "./assets") - + // Fill in the index page rt.GET("/", func(c *gin.Context) { host = &c.Request.Host fillIndex(c) }) + // create a new interactive session + rt.GET("/new", func(c *gin.Context) { + if host == nil { + host = &c.Request.Host + } + + id := uniuri.New() + + c.HTML(http.StatusOK, "term.html", gin.H{ + "title": "interactive terminal", + "path": "/ws_new/" + id, + "id": id, + }) + }) + + rt.GET("/ws_new/:id", func(c *gin.Context) { + id := c.Param("id") + term_conn.ConnectTerm(c.Writer, c.Request, false, id, cmdToExec) + }) + + // create a viewer of an interactive session + rt.GET("/view/:id", func(c *gin.Context) { + id := c.Param("id") + c.HTML(http.StatusOK, "term.html", gin.H{ + "title": "viewer terminal", + "path": "/ws_view/" + id, + }) + }) + + rt.GET("/ws_view/:id", func(c *gin.Context) { + id := c.Param("id") + term_conn.ConnectTerm(c.Writer, c.Request, true, id, nil) + }) + + // start/stop recording the session + rt.GET("/record/:id", func(c *gin.Context) { + id := c.Param("id") + term_conn.StartRecord(id) + }) + + rt.GET("/stop/:id", func(c *gin.Context) { + id := c.Param("id") + term_conn.StopRecord(id) + }) + + // handle static files + rt.Static("/assets", "./assets") + term_conn.Init(checkOrigin) rt.RunTLS(":8080", "./tls/cert.pem", "./tls/private-key.pem") diff --git a/term_conn/reg.go b/term_conn/reg.go index b313b3b..a45bd53 100644 --- a/term_conn/reg.go +++ b/term_conn/reg.go @@ -52,7 +52,20 @@ func (d *Registry) sendToPlayer(name string, ws *websocket.Conn) bool { tc, ok := d.players[name] if ok { - tc.vchan <- ws + tc.viewChan <- ws + } + + d.mtx.Unlock() + return ok +} + +// Send a command to the session to start, stop recording +func (d *Registry) recordSession(id string, cmd int) bool { + d.mtx.Lock() + tc, ok := d.players[id] + + if ok { + tc.recordChan <- cmd } d.mtx.Unlock() diff --git a/term_conn/relay.go b/term_conn/relay.go index 30ce908..b60fbfd 100644 --- a/term_conn/relay.go +++ b/term_conn/relay.go @@ -2,15 +2,16 @@ package term_conn import ( + "encoding/json" "log" "net/http" "os" "os/exec" + "strconv" "sync" "time" "github.com/creack/pty" - "github.com/dchest/uniuri" "github.com/gorilla/websocket" ) @@ -30,6 +31,9 @@ const ( // Send pings to peer with this period. Must be less than pongWait. pingPeriod = (pongWait * 9) / 10 + + recordCmd = 1 + stopCmd = 0 ) var upgrader = websocket.Upgrader{ @@ -46,12 +50,20 @@ type TermConn struct { Name string Ip string - ws *websocket.Conn - ptmx *os.File // the pty that runs the command - cmd *exec.Cmd // represents the process, we need it to terminate the process - vchan chan *websocket.Conn // channel to receive viewers - ws_done chan struct{} // ws is closed, only close this chan in ws reader - pty_done chan struct{} // pty is closed, close this chan in pty reader + ws *websocket.Conn + ptmx *os.File // the pty that runs the command + record *os.File // record session + lastRecTime time.Time // last time a record is written + cmd *exec.Cmd // represents the process, we need it to terminate the process + viewChan chan *websocket.Conn // channel to receive viewers + recordChan chan int // channel to start/stop recording + ws_done chan struct{} // ws is closed, only close this chan in ws reader + pty_done chan struct{} // pty is closed, close this chan in pty reader +} + +type writeRecord struct { + Dur time.Duration `json:"Duration"` + Data []byte `json:"Data"` } func (tc *TermConn) createPty(cmdline []string) error { @@ -230,7 +242,37 @@ out: } } - case viewer := <-tc.vchan: + // Do we need to record the session? + if tc.record != nil { + jbuf, err := json.Marshal(writeRecord{Dur: time.Since(tc.lastRecTime), Data: buf}) + if err != nil { + log.Println("Failed to marshal record", err) + } else { + tc.record.Write(jbuf) + } + + tc.lastRecTime = time.Now() + } + + case cmd := <-tc.recordChan: + var err error + if cmd == recordCmd { + // use the session ID and current as file name + fname := "./records/" + tc.Name + "_" + strconv.FormatInt(time.Now().Unix(), 16) + ".rec" + + tc.record, err = os.OpenFile(fname, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0644) + if err != nil { + log.Println("Failed to create record file", fname, err) + tc.record = nil + } + + tc.lastRecTime = time.Now() + } else { + tc.record.Close() + tc.record = nil + } + + case viewer := <-tc.viewChan: log.Println("Received viewer", viewer.RemoteAddr().String()) viewers = append(viewers, viewer) @@ -287,14 +329,15 @@ func (tc *TermConn) release() { log.Printf("Failed to wait for shell process(%v): %v", proc.Pid, err) } - close(tc.vchan) + close(tc.viewChan) + close(tc.recordChan) } tc.ws.Close() } // handle websockets -func handlePlayer(w http.ResponseWriter, r *http.Request, cmdline []string) { +func handlePlayer(w http.ResponseWriter, r *http.Request, name string, cmdline []string) { ws, err := upgrader.Upgrade(w, r, nil) if err != nil { @@ -304,7 +347,7 @@ func handlePlayer(w http.ResponseWriter, r *http.Request, cmdline []string) { tc := TermConn{ ws: ws, - Name: uniuri.New(), + Name: name, Ip: ws.RemoteAddr().String(), } @@ -313,7 +356,8 @@ func handlePlayer(w http.ResponseWriter, r *http.Request, cmdline []string) { tc.ws_done = make(chan struct{}) tc.pty_done = make(chan struct{}) - tc.vchan = make(chan *websocket.Conn) + tc.viewChan = make(chan *websocket.Conn) + tc.recordChan = make(chan int) if err := tc.createPty(cmdline); err != nil { log.Println("Failed to create PTY: ", err) @@ -353,11 +397,11 @@ func handleViewer(w http.ResponseWriter, r *http.Request, path string) { } } -func ConnectTerm(w http.ResponseWriter, r *http.Request, isViewer bool, path string, cmdline []string) { +func ConnectTerm(w http.ResponseWriter, r *http.Request, isViewer bool, name string, cmdline []string) { if !isViewer { - handlePlayer(w, r, cmdline) + handlePlayer(w, r, name, cmdline) } else { - handleViewer(w, r, path) + handleViewer(w, r, name) } } @@ -365,3 +409,11 @@ func Init(checkOrigin func(r *http.Request) bool) { upgrader.CheckOrigin = checkOrigin registry.init() } + +func StartRecord(id string) { + registry.recordSession(id, recordCmd) +} + +func StopRecord(id string) { + registry.recordSession(id, stopCmd) +}