add record/stop

This commit is contained in:
Zhi Wang 2022-01-14 13:35:10 -05:00
parent ab3a924c3a
commit 8e22acdded
6 changed files with 169 additions and 68 deletions

View File

@ -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>

View File

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

View File

@ -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"
class="d-inline-block align-text-top">
<img src="/assets/logo.svg" style="margin-right: 0.5rem;" height="28" 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}}");
// 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('');
function recordOnOff(on) {
var btn = document.getElementById("record_onoff");
if (btn.value == "Record") {
btn.value = "Stop";
btn.innerHTML = btn.value
fetch("/record/{{.id}}")
} else {
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>

93
main.go
View File

@ -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")

View File

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

View File

@ -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)
}