add record/stop

assets/index.html

@@ -42,9 +42,9 @@
         <div class="card shadow-sm border-info  bg-light mb-3" style="width: 16rem; margin:1em;">
           <div class="card-body">
             <h5 class="card-title">Interactive session</h5>
-            <p class="card-text">From <em>{{.Ip}}</em>, running <strong>{{.Cmd}}</strong>, session ID: <u>{{.Name}}</u>
+            <p class="card-text">From <em>{{.Ip}}</em>, running <strong>{{.Cmd}}</strong>, session ID: <u>{{.Id}}</u>
             </p>
-            <a class="btn btn-secondary btn-sm float-end" href="/view/{{.Name}}" target="_blank" role="button">View
+            <a class="btn btn-secondary btn-sm float-end" href="/view/{{.Id}}" target="_blank" role="button">View
               Session</a>
           </div>
         </div>

assets/main.css

@@ -35,4 +35,4 @@
 .navbar-xs .navbar-nav>li>a {
     padding-top: 0px;
     padding-bottom: 0px;
-}
+}

assets/term.html

@@ -23,11 +23,11 @@
     <nav class="navbar navbar-dark bg-dark shadow-sm navbar-xs">
       <div class="container-fluid">
         <a class="navbar-brand mx-auto" href="https://github.com/syssecfsu/witty" target="_blank">
-          <img src="/assets/logo.svg" style="margin-right: 0.5rem;" height="28"
+          <img src="/assets/logo.svg" style="margin-right: 0.5rem;" height="28" class="d-inline-block align-text-top">
-            class="d-inline-block align-text-top">
           {{.title}}
         </a>
-        <!-- <a class="btn btn-primary btn-sm float-end" href="/new" role="button">New Session</a> -->
+        <button type="button" id="record_onoff" class="btn btn-primary btn-sm float-end" value="Record"
+          onclick="recordOnOff()">Record</button>
       </div>
     </nav>
   </header>
@@ -40,16 +40,33 @@
   </div>
 
   <script>
-    term = createTerminal("{{.path}}");
+    function recordOnOff(on) {
-    // print something to test output and scroll
+      var btn = document.getElementById("record_onoff");
-    var str = [
+      if (btn.value == "Record") {
-      ' ┌────────────────────────────────────────────────────────────────────────────┐\n',
+        btn.value = "Stop";
-      ' │                       Powered by \u001b[32;1mGo, Gin, websocket, pty, and xterm.js\x1b[0m     │\n',
+        btn.innerHTML = btn.value
-      ' └────────────────────────────────────────────────────────────────────────────┘\n',
+        fetch("/record/{{.id}}")
-      ''
+      } else {
-    ].join('');
+        btn.value = "Record"; 
+        btn.innerHTML = btn.value
+        fetch("/stop/{{.id}}")
+      }
+    }
 
-    term.writeln(str);
+    function Init() {
+      term = createTerminal("{{.path}}");
+      // print something to test output and scroll
+      var str = [
+        ' ┌────────────────────────────────────────────────────────────────────────────┐\n',
+        ' │                       Powered by \u001b[32;1mGo, Gin, websocket, pty, and xterm.js\x1b[0m     │\n',
+        ' └────────────────────────────────────────────────────────────────────────────┘\n',
+        ''
+      ].join('');
+
+      term.writeln(str);
+    }
+
+    Init()
   </script>
 
 </body>

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
+	Ip  string
-	Cmd  string
+	Cmd string
-	Name 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,
+			Id:  tc.Name,
-			Ip:   tc.Ip,
+			Ip:  tc.Ip,
-			Cmd:  cmdToExec[0],
+			Cmd: cmdToExec[0],
 		})
 	})
 
@@ -79,42 +80,60 @@ func main() {
 	rt.SetTrustedProxies(nil)
 	rt.LoadHTMLGlob("./assets/*.html")
 
-	rt.GET("/view/:sname", func(c *gin.Context) {
+	// Fill in the index page
-		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")
-
 	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")

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()

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
+	ws          *websocket.Conn
-	ptmx     *os.File             // the pty that runs the command
+	ptmx        *os.File             // the pty that runs the command
-	cmd      *exec.Cmd            // represents the process, we need it to terminate the process
+	record      *os.File             // record session
-	vchan    chan *websocket.Conn // channel to receive viewers
+	lastRecTime time.Time            // last time a record is written
-	ws_done  chan struct{}        // ws is closed, only close this chan in ws reader
+	cmd         *exec.Cmd            // represents the process, we need it to terminate the process
-	pty_done chan struct{}        // pty is closed, close this chan in pty reader
+	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)
+}