From 05c0f4188be0a440fd8c3514fe55e5bba5cfea61 Mon Sep 17 00:00:00 2001
From: Jason <xjasonlyu@gmail.com>
Date: Sun, 9 Aug 2020 13:59:08 +0800
Subject: [PATCH] Add json status

---
 component/session/server.go  | 63 ++++++++++++++++++++++++++++--------
 component/session/session.go | 32 ++++++++++++------
 2 files changed, 72 insertions(+), 23 deletions(-)

diff --git a/component/session/server.go b/component/session/server.go
index c0a7405..5a52319 100644
--- a/component/session/server.go
+++ b/component/session/server.go
@@ -2,6 +2,7 @@ package session
 
 import (
 	"bufio"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -18,7 +19,7 @@ import (
 	"github.com/gobuffalo/packr"
 )
 
-const maxCompletedSessions = 100
+const maxClosedSessions = 100
 
 type Server struct {
 	sync.Mutex
@@ -30,7 +31,7 @@ type Server struct {
 	trafficDown int64
 
 	activeSessionMap  sync.Map
-	completedSessions []Session
+	closedSessionList []Session
 }
 
 func New(addr string) *Server {
@@ -39,19 +40,52 @@ func New(addr string) *Server {
 	}
 }
 
-func (s *Server) handler(resp http.ResponseWriter, req *http.Request) {
+func (s *Server) getSessions() (activeSessions, closedSessions []Session) {
 	// Slice of active sessions
-	var activeSessions []Session
 	s.activeSessionMap.Range(func(key, value interface{}) bool {
 		session := value.(*Session)
 		activeSessions = append(activeSessions, *session)
 		return true
 	})
 
-	// Slice of completed sessions
+	// Slice of closed sessions
 	s.Lock()
-	completedSessions := append([]Session(nil), s.completedSessions...)
-	s.Unlock()
+	defer s.Unlock()
+	closedSessions = append([]Session(nil), s.closedSessionList...)
+	return
+}
+
+func (s *Server) serveJSON(w http.ResponseWriter, _ *http.Request) {
+	activeSessions, closedSessions := s.getSessions()
+
+	// calculate traffic
+	trafficUp := atomic.LoadInt64(&s.trafficUp)
+	trafficDown := atomic.LoadInt64(&s.trafficDown)
+	for _, session := range activeSessions {
+		trafficUp += session.UploadBytes
+		trafficDown += session.DownloadBytes
+	}
+
+	status := &Status{
+		platform(),
+		C.Version,
+		cpu(),
+		mem(),
+		uptime(),
+		trafficUp + trafficDown,
+		trafficUp,
+		trafficDown,
+		runtime.NumGoroutine(),
+		activeSessions,
+		closedSessions,
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(status)
+}
+
+func (s *Server) serveHTML(resp http.ResponseWriter, _ *http.Request) {
+	activeSessions, closedSessions := s.getSessions()
 
 	tablePrint := func(w io.Writer, sessions []Session) {
 		// Sort by session start time.
@@ -121,8 +155,8 @@ func (s *Server) handler(resp http.ResponseWriter, req *http.Request) {
 	// Session table
 	_, _ = fmt.Fprintf(w, "<h3 class=\"sub-header\">Active sessions (%d)</h3>\n", len(activeSessions))
 	tablePrint(w, activeSessions)
-	_, _ = fmt.Fprintf(w, "<h3 class=\"sub-header\">Closed sessions (%d)</h3>\n", len(completedSessions))
-	tablePrint(w, completedSessions)
+	_, _ = fmt.Fprintf(w, "<h3 class=\"sub-header\">Closed sessions (%d)</h3>\n", len(closedSessions))
+	tablePrint(w, closedSessions)
 	_, _ = fmt.Fprintf(w, "</div></body></html>\n")
 	_ = w.Flush()
 }
@@ -144,7 +178,8 @@ func (s *Server) Start() error {
 	}
 
 	mux := http.NewServeMux()
-	mux.HandleFunc("/", s.handler)
+	mux.HandleFunc("/", s.serveHTML)
+	mux.HandleFunc("/json", s.serveJSON)
 
 	box := packr.NewBox("./css")
 	mux.Handle("/css/", http.StripPrefix("/css/", http.FileServer(box)))
@@ -175,11 +210,11 @@ func (s *Server) RemoveSession(key interface{}) {
 		// record up & down traffic
 		atomic.AddInt64(&s.trafficUp, atomic.LoadInt64(&session.UploadBytes))
 		atomic.AddInt64(&s.trafficDown, atomic.LoadInt64(&session.DownloadBytes))
-		// move to completed sessions
+		// move to closed sessions
 		s.Lock()
-		s.completedSessions = append(s.completedSessions, *session)
-		if len(s.completedSessions) > maxCompletedSessions {
-			s.completedSessions = s.completedSessions[1:]
+		s.closedSessionList = append(s.closedSessionList, *session)
+		if len(s.closedSessionList) > maxClosedSessions {
+			s.closedSessionList = s.closedSessionList[1:]
 		}
 		s.Unlock()
 	}
diff --git a/component/session/session.go b/component/session/session.go
index 778c5ed..8ce0aef 100644
--- a/component/session/session.go
+++ b/component/session/session.go
@@ -16,16 +16,30 @@ type Monitor interface {
 	RemoveSession(key interface{})
 }
 
+type Status struct {
+	Platform       string    `json:"platform"`
+	Version        string    `json:"version"`
+	CPU            string    `json:"cpu"`
+	MEM            string    `json:"mem"`
+	Uptime         string    `json:"uptime"`
+	Total          int64     `json:"total"`
+	Upload         int64     `json:"upload"`
+	Download       int64     `json:"download"`
+	Goroutines     int       `json:"goroutines"`
+	ActiveSessions []Session `json:"activeSessions"`
+	ClosedSessions []Session `json:"closedSessions"`
+}
+
 type Session struct {
-	Process       string
-	Network       string
-	DialerAddr    string
-	ClientAddr    string
-	TargetAddr    string
-	UploadBytes   int64
-	DownloadBytes int64
-	SessionStart  time.Time
-	SessionClose  time.Time
+	Process       string    `json:"process"`
+	Network       string    `json:"network"`
+	DialerAddr    string    `json:"dialerAddr"`
+	ClientAddr    string    `json:"clientAddr"`
+	TargetAddr    string    `json:"targetAddr"`
+	UploadBytes   int64     `json:"upload"`
+	DownloadBytes int64     `json:"download"`
+	SessionStart  time.Time `json:"sessionStart"`
+	SessionClose  time.Time `json:"sessionClose"`
 }
 
 // Track SessionConn