diff options
Diffstat (limited to 'internal')
| -rw-r--r-- | internal/ipc/ipc.go | 26 | ||||
| -rw-r--r-- | internal/ipc/server.go | 22 | ||||
| -rw-r--r-- | internal/store/store.go | 9 |
3 files changed, 50 insertions, 7 deletions
diff --git a/internal/ipc/ipc.go b/internal/ipc/ipc.go index 03a4bc7..d295238 100644 --- a/internal/ipc/ipc.go +++ b/internal/ipc/ipc.go @@ -48,11 +48,12 @@ type Request struct { // Response is sent by the daemon to a client. type Response struct { - Type string `json:"type"` // "list", "get", "new", "repeat", "rules", "intrude_result", "intrude_done", or "error" + Type string `json:"type"` // "list", "get", "new", "repeat", "rules", "intrude_result", "intrude_done", "status", or "error" Entries []store.Summary `json:"entries,omitempty"` // for "list" Detail *EntryDetail `json:"detail,omitempty"` // for "get" and "repeat" New *store.Summary `json:"new,omitempty"` // for "new" (subscribe push) Rules []rules.Rule `json:"rules,omitempty"` // for "rules" + Status *StatusMsg `json:"status,omitempty"` // for "status" // For "intrude_result": one completed attack request. IntrudeResult *IntrudeResultMsg `json:"intrude_result,omitempty"` @@ -60,6 +61,12 @@ type Response struct { Error string `json:"error,omitempty"` } +// StatusMsg is basic daemon info for a TUI status bar. +type StatusMsg struct { + ProxyAddr string `json:"proxy_addr"` + HistoryCount int64 `json:"history_count"` +} + // IntrudeResultMsg is one completed Intruder attack request. type IntrudeResultMsg struct { Position int `json:"position"` @@ -110,6 +117,23 @@ func (c *Client) Close() error { return c.conn.Close() } +// Status returns basic daemon info for a status bar. +func (c *Client) Status() (*StatusMsg, error) { + c.mu.Lock() + defer c.mu.Unlock() + if err := c.enc.Encode(Request{Type: "status"}); err != nil { + return nil, err + } + var resp Response + if err := c.dec.Decode(&resp); err != nil { + return nil, err + } + if resp.Type == "error" { + return nil, errors.New(resp.Error) + } + return resp.Status, nil +} + // List returns up to limit history summaries older than beforeID (0 for // the most recent), newest first. func (c *Client) List(limit int, beforeID int64) ([]store.Summary, error) { diff --git a/internal/ipc/server.go b/internal/ipc/server.go index 50dff7f..eea0ebb 100644 --- a/internal/ipc/server.go +++ b/internal/ipc/server.go @@ -66,16 +66,18 @@ func (h *Hub) unsubscribe(ch chan store.Summary) { // Server serves the daemon side of the mitmux control protocol. type Server struct { - db *store.Store - hub *Hub - repeater Repeater - intruder Intruder + db *store.Store + hub *Hub + repeater Repeater + intruder Intruder + proxyAddr string } // NewServer creates a control-protocol Server backed by db, broadcasting // through hub and sending Repeater/Intruder requests through rep. -func NewServer(db *store.Store, hub *Hub, rep Repeater) *Server { - s := &Server{db: db, hub: hub, repeater: rep} +// proxyAddr is reported back to clients via the "status" request. +func NewServer(db *store.Store, hub *Hub, rep Repeater, proxyAddr string) *Server { + s := &Server{db: db, hub: hub, repeater: rep, proxyAddr: proxyAddr} if in, ok := rep.(Intruder); ok { s.intruder = in } @@ -106,6 +108,14 @@ func (s *Server) handleConn(conn net.Conn) { } switch req.Type { + case "status": + count, err := s.db.Count() + if err != nil { + enc.Encode(Response{Type: "error", Error: err.Error()}) + continue + } + enc.Encode(Response{Type: "status", Status: &StatusMsg{ProxyAddr: s.proxyAddr, HistoryCount: count}}) + case "list": var entries []store.Summary var err error diff --git a/internal/store/store.go b/internal/store/store.go index 211162e..eb017db 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -185,6 +185,15 @@ func (s *Store) Insert(e *Entry) (int64, error) { return id, nil } +// Count returns the total number of history entries. +func (s *Store) Count() (int64, error) { + var n int64 + if err := s.db.QueryRow(`SELECT count(*) FROM history`).Scan(&n); err != nil { + return 0, fmt.Errorf("count history: %w", err) + } + return n, nil +} + // List returns up to limit history summaries older than beforeID (or the // most recent if beforeID is 0), newest first. func (s *Store) List(limit int, beforeID int64) ([]Summary, error) { |