diff options
Diffstat (limited to 'internal/store')
| -rw-r--r-- | internal/store/annotation.go | 58 | ||||
| -rw-r--r-- | internal/store/annotation_test.go | 92 | ||||
| -rw-r--r-- | internal/store/har.go | 239 | ||||
| -rw-r--r-- | internal/store/history.go | 401 | ||||
| -rw-r--r-- | internal/store/history_test.go | 190 | ||||
| -rw-r--r-- | internal/store/intercept.go | 233 | ||||
| -rw-r--r-- | internal/store/intruder.go | 72 | ||||
| -rw-r--r-- | internal/store/loot.go | 103 | ||||
| -rw-r--r-- | internal/store/loot_test.go | 155 | ||||
| -rw-r--r-- | internal/store/plugin.go | 156 | ||||
| -rw-r--r-- | internal/store/plugin_test.go | 368 | ||||
| -rw-r--r-- | internal/store/project.go | 245 | ||||
| -rw-r--r-- | internal/store/repeater.go | 103 | ||||
| -rw-r--r-- | internal/store/repeater_test.go | 265 | ||||
| -rw-r--r-- | internal/store/scope.go | 116 | ||||
| -rw-r--r-- | internal/store/scope_test.go | 135 | ||||
| -rw-r--r-- | internal/store/search_test.go | 105 | ||||
| -rw-r--r-- | internal/store/store_test.go | 22 | ||||
| -rw-r--r-- | internal/store/websocket.go | 198 | ||||
| -rw-r--r-- | internal/store/websocket_test.go | 200 |
20 files changed, 3456 insertions, 0 deletions
diff --git a/internal/store/annotation.go b/internal/store/annotation.go new file mode 100644 index 0000000..84b0c48 --- /dev/null +++ b/internal/store/annotation.go @@ -0,0 +1,58 @@ +package store + +import "fmt" + +// Annotation holds optional comment and highlight colour for a history row. +type Annotation struct { + HistoryID uint64 + Comment string + Color string // empty means no highlight; one of: "red", "orange", "yellow", "green", "blue", "purple" +} + +// SetAnnotation creates or replaces the annotation for the given history row. +func (s *Store) SetAnnotation(a Annotation) error { + return s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO annotations (history_id, comment, color) + VALUES (?, ?, ?) + ON CONFLICT(history_id) DO UPDATE SET + comment = excluded.comment, + color = excluded.color`, + a.HistoryID, a.Comment, a.Color, + ) + if err != nil { + return fmt.Errorf("store: set annotation %d: %w", a.HistoryID, err) + } + return nil + }) +} + +// DeleteAnnotation removes the annotation for the given history row. +func (s *Store) DeleteAnnotation(historyID uint64) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM annotations WHERE history_id = ?`, historyID) + if err != nil { + return fmt.Errorf("store: delete annotation %d: %w", historyID, err) + } + return nil + }) +} + +// AllAnnotations returns every annotation row. Called once at startup to +// populate the in-memory annotation map. +func (s *Store) AllAnnotations() ([]Annotation, error) { + rows, err := s.db.Query(`SELECT history_id, comment, color FROM annotations`) + if err != nil { + return nil, fmt.Errorf("store: all annotations: %w", err) + } + defer rows.Close() + var out []Annotation + for rows.Next() { + var a Annotation + if err := rows.Scan(&a.HistoryID, &a.Comment, &a.Color); err != nil { + return nil, fmt.Errorf("store: scan annotation: %w", err) + } + out = append(out, a) + } + return out, rows.Err() +} diff --git a/internal/store/annotation_test.go b/internal/store/annotation_test.go new file mode 100644 index 0000000..72c79e1 --- /dev/null +++ b/internal/store/annotation_test.go @@ -0,0 +1,92 @@ +package store_test + +import ( + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func logTxForAnnotation(t *testing.T, st *store.Store) uint64 { + t.Helper() + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + txs, err := st.AllTransactions() + require.NoError(t, err) + require.NotEmpty(t, txs) + return txs[0].ID +} + +func TestSetAnnotation_PersistsCommentAndColor(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + a := store.Annotation{HistoryID: id, Comment: "suspicious", Color: "red"} + require.NoError(t, st.SetAnnotation(a)) + + all, err := st.AllAnnotations() + require.NoError(t, err) + require.Len(t, all, 1) + assert.Equal(t, id, all[0].HistoryID) + assert.Equal(t, "suspicious", all[0].Comment) + assert.Equal(t, "red", all[0].Color) +} + +func TestSetAnnotation_UpdatesExisting(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "first", Color: "blue"})) + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "second", Color: "green"})) + + all, err := st.AllAnnotations() + require.NoError(t, err) + require.Len(t, all, 1) + assert.Equal(t, "second", all[0].Comment) + assert.Equal(t, "green", all[0].Color) +} + +func TestDeleteAnnotation_RemovesRow(t *testing.T) { + st := newTestStore(t) + id := logTxForAnnotation(t, st) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: id, Comment: "note", Color: "yellow"})) + require.NoError(t, st.DeleteAnnotation(id)) + + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Empty(t, all) +} + +func TestDeleteAnnotation_NonExistent_NoError(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.DeleteAnnotation(9999)) +} + +func TestAllAnnotations_Empty(t *testing.T) { + st := newTestStore(t) + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Empty(t, all) +} + +func TestAllAnnotations_MultipleRows(t *testing.T) { + st := newTestStore(t) + + tx1 := makeTx("GET", "a.com", "https://a.com/", 200) + tx2 := makeTx("POST", "b.com", "https://b.com/", 201) + require.NoError(t, st.Log(tx1)) + require.NoError(t, st.Log(tx2)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 2) + + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: txs[0].ID, Comment: "c1", Color: "red"})) + require.NoError(t, st.SetAnnotation(store.Annotation{HistoryID: txs[1].ID, Comment: "c2", Color: "blue"})) + + all, err := st.AllAnnotations() + require.NoError(t, err) + assert.Len(t, all, 2) +} diff --git a/internal/store/har.go b/internal/store/har.go new file mode 100644 index 0000000..047300a --- /dev/null +++ b/internal/store/har.go @@ -0,0 +1,239 @@ +package store + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "net/url" + "strings" + "time" +) + +// HAR 1.2 — https://w3c.github.io/web-performance/specs/HAR/Overview.html + +type HAR struct { + Log HARLog `json:"log"` +} + +type HARLog struct { + Version string `json:"version"` + Creator HARCreator `json:"creator"` + Entries []HAREntry `json:"entries"` +} + +type HARCreator struct { + Name string `json:"name"` + Version string `json:"version"` +} + +type HAREntry struct { + StartedDateTime string `json:"startedDateTime"` + Time int64 `json:"time"` + Request HARRequest `json:"request"` + Response HARResponse `json:"response"` + Timings HARTimings `json:"timings"` +} + +type HARRequest struct { + Method string `json:"method"` + URL string `json:"url"` + HTTPVersion string `json:"httpVersion"` + Headers []HARNameValue `json:"headers"` + QueryString []HARNameValue `json:"queryString"` + PostData *HARPostData `json:"postData,omitempty"` + HeadersSize int `json:"headersSize"` + BodySize int `json:"bodySize"` +} + +type HARResponse struct { + Status int `json:"status"` + StatusText string `json:"statusText"` + HTTPVersion string `json:"httpVersion"` + Headers []HARNameValue `json:"headers"` + Content HARContent `json:"content"` + RedirectURL string `json:"redirectURL"` + HeadersSize int `json:"headersSize"` + BodySize int `json:"bodySize"` +} + +type HARNameValue struct { + Name string `json:"name"` + Value string `json:"value"` +} + +type HARPostData struct { + MimeType string `json:"mimeType"` + Text string `json:"text"` +} + +type HARContent struct { + Size int `json:"size"` + MimeType string `json:"mimeType"` + Text string `json:"text,omitempty"` +} + +type HARTimings struct { + Send int64 `json:"send"` + Wait int64 `json:"wait"` + Receive int64 `json:"receive"` +} + +// ExportHAR serialises the given transactions into a HAR 1.2 JSON byte slice. +// Pass only the transactions you want exported — scope and out-of-scope +// filtering is the caller's responsibility. +func ExportHAR(txs []Transaction) ([]byte, error) { + har := HAR{ + Log: HARLog{ + Version: "1.2", + Creator: HARCreator{ + Name: "Shiv", + Version: "1.0", + }, + Entries: make([]HAREntry, 0, len(txs)), + }, + } + + for _, tx := range txs { + entry := HAREntry{ + StartedDateTime: tx.Timestamp.UTC().Format(time.RFC3339), + Time: tx.DurationMs, + Request: buildHARRequest(tx), + Response: buildHARResponse(tx), + Timings: HARTimings{ + Send: 0, + Wait: tx.DurationMs, + Receive: 0, + }, + } + har.Log.Entries = append(har.Log.Entries, entry) + } + + data, err := json.MarshalIndent(har, "", " ") + if err != nil { + return nil, fmt.Errorf("har: marshal: %w", err) + } + return data, nil +} + +func stripDefaultPort(rawURL string) string { + for _, prefix := range []string{"https://", "http://"} { + if !strings.HasPrefix(rawURL, prefix) { + continue + } + rest := rawURL[len(prefix):] + slash := strings.Index(rest, "/") + if slash < 0 { + slash = len(rest) + } + host := rest[:slash] + path := rest[slash:] + h, port, err := net.SplitHostPort(host) + if err != nil { + return rawURL + } + if (prefix == "https://" && port == "443") || (prefix == "http://" && port == "80") { + return prefix + h + path + } + return rawURL + } + return rawURL +} + +func buildHARRequest(tx Transaction) HARRequest { + headers := make([]HARNameValue, 0, len(tx.ReqHeaders)) + for k, vals := range tx.ReqHeaders { + for _, v := range vals { + headers = append(headers, HARNameValue{Name: k, Value: v}) + } + } + + httpVersion := tx.Proto + if httpVersion == "" { + httpVersion = "HTTP/1.1" + } + + req := HARRequest{ + Method: tx.Method, + URL: stripDefaultPort(tx.URL), + HTTPVersion: httpVersion, + Headers: headers, + QueryString: extractQueryString(tx.URL), + HeadersSize: -1, + BodySize: len(tx.ReqBody), + } + + if len(tx.ReqBody) > 0 { + mimeType := tx.ReqHeaders.Get("Content-Type") + if mimeType == "" { + mimeType = "application/octet-stream" + } + req.PostData = &HARPostData{ + MimeType: mimeType, + Text: string(tx.ReqBody), + } + } + + return req +} + +func buildHARResponse(tx Transaction) HARResponse { + headers := make([]HARNameValue, 0, len(tx.RespHeaders)) + for k, vals := range tx.RespHeaders { + for _, v := range vals { + headers = append(headers, HARNameValue{Name: k, Value: v}) + } + } + + httpVersion := tx.Proto + if httpVersion == "" { + httpVersion = "HTTP/1.1" + } + + mimeType := tx.RespHeaders.Get("Content-Type") + if mimeType == "" { + mimeType = "application/octet-stream" + } + + redirectURL := tx.RespHeaders.Get("Location") + + return HARResponse{ + Status: tx.StatusCode, + StatusText: statusText(tx.StatusCode), + HTTPVersion: httpVersion, + Headers: headers, + Content: HARContent{ + Size: len(tx.RespBody), + MimeType: mimeType, + Text: string(tx.RespBody), + }, + RedirectURL: redirectURL, + HeadersSize: -1, + BodySize: len(tx.RespBody), + } +} + +func extractQueryString(rawURL string) []HARNameValue { + idx := strings.Index(rawURL, "?") + if idx < 0 { + return []HARNameValue{} + } + parsed, err := url.ParseQuery(rawURL[idx+1:]) + if err != nil || len(parsed) == 0 { + return []HARNameValue{} + } + pairs := make([]HARNameValue, 0, len(parsed)) + for k, vals := range parsed { + for _, v := range vals { + pairs = append(pairs, HARNameValue{Name: k, Value: v}) + } + } + return pairs +} + +func statusText(code int) string { + if t := http.StatusText(code); t != "" { + return t + } + return fmt.Sprintf("%d", code) +} diff --git a/internal/store/history.go b/internal/store/history.go new file mode 100644 index 0000000..066c7c2 --- /dev/null +++ b/internal/store/history.go @@ -0,0 +1,401 @@ +package store + +import ( + "encoding/json" + "fmt" + "net" + "net/http" + "strconv" + "strings" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +// Transaction is a matched HTTP request/response pair. +// Host is always a pure hostname — never host:port. +// Port holds the destination port as an integer. +type Transaction struct { + ID uint64 + Timestamp time.Time + Host string + Port int + Method string + URL string + Proto string // "HTTP/2" or "HTTP/1.1" + ReqHeaders http.Header + ReqBody []byte + StatusCode int + RespHeaders http.Header + RespBody []byte + RespSize int + DurationMs int64 + TLS bool + InScope bool +} + +// TransactionFilter holds filter criteria for paginated history queries. +// Zero values mean "no filter" for that field. +type TransactionFilter struct { + Search string // matched against host + url + method + status (case-insensitive) + Host string // exact host match (from site map selection) + Port int // exact port match; 0 means no filter + PathPrefix string // URL path prefix filter (from site map selection) + ShowOutScope bool // if false, only in_scope=1 rows are returned + Methods []string // whitelist of HTTP methods; empty = all + StatusMin int // lower bound on status_code; 0 = no lower bound + StatusMax int // upper bound on status_code; 0 = no upper bound + HideStatic bool // if true, filter out common static asset extensions in memory +} + +// PageSize is the number of rows returned per page. +const PageSize = 100 + +// SiteMapEntry is a minimal host+port+url triplet used to build the site map +// at startup. Port is required to differentiate apps on the same host at +// different ports (e.g. localhost:3000 vs localhost:8080). +type SiteMapEntry struct { + Host string + Port int + URL string +} + +// SiteMapEntries returns the host, port, and url for every row in history. +// Used once at startup to populate the site map with all historical +// paths before the first page of transactions is displayed. +func (s *Store) SiteMapEntries() ([]SiteMapEntry, error) { + rows, err := s.db.Query(`SELECT host, port, url FROM history ORDER BY id ASC`) + if err != nil { + return nil, fmt.Errorf("store: site map entries: %w", err) + } + defer rows.Close() + var entries []SiteMapEntry + for rows.Next() { + var e SiteMapEntry + if err := rows.Scan(&e.Host, &e.Port, &e.URL); err != nil { + return nil, fmt.Errorf("store: scan site map entry: %w", err) + } + entries = append(entries, e) + } + return entries, rows.Err() +} + +func (s *Store) GetTransaction(id uint64) (*Transaction, error) { + var transaction Transaction + var timestampStr, reqHeaders, respHeaders string + var tlsFlag, scopeFlag int + err := s.db.QueryRow(` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, req_body, status_code, resp_headers, + resp_body, duration_ms, tls, in_scope, resp_size + FROM history WHERE id = ?`, id, + ).Scan( + &transaction.ID, ×tampStr, &transaction.Host, &transaction.Port, + &transaction.Method, &transaction.URL, &transaction.Proto, + &reqHeaders, &transaction.ReqBody, &transaction.StatusCode, &respHeaders, + &transaction.RespBody, &transaction.DurationMs, &tlsFlag, &scopeFlag, + &transaction.RespSize, + ) + if err != nil { + return nil, fmt.Errorf("store: get transaction %d: %w", id, err) + } + transaction.Timestamp, _ = time.Parse(time.RFC3339, timestampStr) + transaction.TLS = tlsFlag == 1 + transaction.InScope = scopeFlag == 1 + if err := json.Unmarshal([]byte(reqHeaders), &transaction.ReqHeaders); err != nil { + transaction.ReqHeaders = http.Header{} + } + if err := json.Unmarshal([]byte(respHeaders), &transaction.RespHeaders); err != nil { + transaction.RespHeaders = http.Header{} + } + return &transaction, nil +} + +// Log writes every transaction as a unique row — no deduplication. +func (s *Store) Log(t Transaction) error { + return s.write(func() error { + reqH, err := json.Marshal(t.ReqHeaders) + if err != nil { + return fmt.Errorf("store: marshal req headers: %w", err) + } + respH, err := json.Marshal(t.RespHeaders) + if err != nil { + return fmt.Errorf("store: marshal resp headers: %w", err) + } + proto := t.Proto + if proto == "" { + proto = "HTTP/1.1" + } + respSize := len(t.RespBody) + res, err := s.db.Exec(` + INSERT INTO history + (timestamp, host, port, method, url, proto, req_headers, req_body, + status_code, resp_headers, resp_body, duration_ms, tls, in_scope, resp_size) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + t.Timestamp.UTC().Format(time.RFC3339), + t.Host, t.Port, t.Method, t.URL, proto, + string(reqH), t.ReqBody, + t.StatusCode, string(respH), t.RespBody, + t.DurationMs, + boolToInt(t.TLS), + boolToInt(t.InScope), + respSize, + ) + if err != nil { + return fmt.Errorf("store: log transaction: %w", err) + } + id, _ := res.LastInsertId() + t.ID = uint64(id) + t.Proto = proto + t.RespSize = respSize + // Guard against sending on a closed channel when the store is shutting down. + select { + case <-s.done: + default: + select { + case s.Updates <- t: + default: + } + } + return nil + }) +} + +// TransactionsPage returns up to PageSize transactions with id < beforeID, +// ordered by id descending (newest first), applying filter criteria. +// Pass beforeID=0 to start from the newest row. +// Bodies are omitted; use GetTransaction for the full row. +func (s *Store) TransactionsPage(beforeID uint64, filter TransactionFilter) ([]Transaction, error) { + var conditions []string + var args []any + + if beforeID > 0 { + conditions = append(conditions, "id < ?") + args = append(args, beforeID) + } + + if !filter.ShowOutScope { + conditions = append(conditions, "in_scope = 1") + } + + if filter.Host != "" { + conditions = append(conditions, "host = ?") + args = append(args, filter.Host) + if filter.Port != 0 { + conditions = append(conditions, "port = ?") + args = append(args, filter.Port) + } + } + + if filter.PathPrefix != "" { + conditions = append(conditions, "(url LIKE ? OR url LIKE ? OR url LIKE ?)") + prefix := filter.PathPrefix + args = append(args, + "%"+prefix, + "%"+prefix+"?%", + "%"+prefix+"/%", + ) + } + + if filter.Search != "" { + terms := strings.Fields(strings.ToLower(filter.Search)) + for _, term := range terms { + conditions = append(conditions, "(LOWER(host) LIKE ? OR LOWER(url) LIKE ? OR LOWER(method) LIKE ? OR CAST(status_code AS TEXT) LIKE ?)") + like := "%" + term + "%" + args = append(args, like, like, like, like) + } + } + + if len(filter.Methods) > 0 { + placeholders := make([]string, len(filter.Methods)) + for i, m := range filter.Methods { + placeholders[i] = "?" + args = append(args, strings.ToUpper(m)) + } + conditions = append(conditions, "method IN ("+strings.Join(placeholders, ",")+")") + } + + if filter.StatusMin > 0 { + conditions = append(conditions, "status_code >= ?") + args = append(args, filter.StatusMin) + } + if filter.StatusMax > 0 { + conditions = append(conditions, "status_code <= ?") + args = append(args, filter.StatusMax) + } + + query := ` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, '' as req_body, status_code, resp_headers, + '' as resp_body, duration_ms, tls, in_scope, resp_size + FROM history` + + if len(conditions) > 0 { + query += " WHERE " + strings.Join(conditions, " AND ") + } + query += fmt.Sprintf(" ORDER BY id DESC LIMIT %d", PageSize) + + rows, err := s.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("store: query history page: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// SearchTransactions performs a full-text LIKE scan across request/response bodies +// and/or headers. At least one of inReqBody, inRespBody, inHeaders must be true. +// Results are ordered by id descending and capped at 200 rows. +func (s *Store) SearchTransactions(term string, inReqBody, inRespBody, inHeaders bool) ([]Transaction, error) { + if term == "" || (!inReqBody && !inRespBody && !inHeaders) { + return nil, nil + } + like := "%" + term + "%" + var clauses []string + var args []any + if inReqBody { + clauses = append(clauses, "CAST(req_body AS TEXT) LIKE ?") + args = append(args, like) + } + if inRespBody { + clauses = append(clauses, "CAST(resp_body AS TEXT) LIKE ?") + args = append(args, like) + } + if inHeaders { + clauses = append(clauses, "req_headers LIKE ?") + args = append(args, like) + clauses = append(clauses, "resp_headers LIKE ?") + args = append(args, like) + } + query := ` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, '' as req_body, status_code, resp_headers, + '' as resp_body, duration_ms, tls, in_scope, resp_size + FROM history + WHERE (` + strings.Join(clauses, " OR ") + `) + ORDER BY id DESC LIMIT 200` + rows, err := s.db.Query(query, args...) + if err != nil { + return nil, fmt.Errorf("store: search transactions: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// AllTransactions returns all transactions ordered by id descending, capped at +// 500 rows. Used primarily in tests; production UI uses TransactionsPage. +func (s *Store) AllTransactions() ([]Transaction, error) { + rows, err := s.db.Query(` + SELECT id, timestamp, host, port, method, url, proto, + req_headers, req_body, status_code, resp_headers, + resp_body, duration_ms, tls, in_scope, resp_size + FROM history + ORDER BY id DESC LIMIT 500`) + if err != nil { + return nil, fmt.Errorf("store: query all transactions: %w", err) + } + defer rows.Close() + return scanTransactions(rows) +} + +// ClearHistory deletes all rows from the history table. +func (s *Store) ClearHistory() error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM history`) + if err != nil { + return fmt.Errorf("store: clear history: %w", err) + } + _, err = s.db.Exec(`VACUUM`) + if err != nil { + return fmt.Errorf("store: VACUUM history: %w", err) + } + + return nil + }) +} + +func scanTransactions(rows interface { + Next() bool + Scan(...any) error + Err() error +}) ([]Transaction, error) { + var txs []Transaction + for rows.Next() { + var tx Transaction + var ts, reqH, respH string + var tlsInt, scopeInt int + if err := rows.Scan( + &tx.ID, &ts, &tx.Host, &tx.Port, &tx.Method, &tx.URL, &tx.Proto, + &reqH, &tx.ReqBody, &tx.StatusCode, &respH, + &tx.RespBody, &tx.DurationMs, &tlsInt, &scopeInt, + &tx.RespSize, + ); err != nil { + return nil, fmt.Errorf("store: scan transaction: %w", err) + } + tx.Timestamp, _ = time.Parse(time.RFC3339, ts) + tx.TLS = tlsInt == 1 + tx.InScope = scopeInt == 1 + if err := json.Unmarshal([]byte(reqH), &tx.ReqHeaders); err != nil { + tx.ReqHeaders = http.Header{} + } + if err := json.Unmarshal([]byte(respH), &tx.RespHeaders); err != nil { + tx.RespHeaders = http.Header{} + } + txs = append(txs, tx) + } + return txs, rows.Err() +} + +// ObserveResponse implements events.ResponseObserver. It splits e.Host +// (which arrives as host:port from the proxy) into a pure hostname and integer +// port once here — no downstream consumer ever parses host/port again. +func (s *Store) ObserveResponse(e events.ResponseEvent) { + host, portStr, err := net.SplitHostPort(e.Host) + if err != nil { + host = e.Host + if e.TLS { + portStr = "443" + } else { + portStr = "80" + } + } + port, _ := strconv.Atoi(portStr) + + if err := s.Log(Transaction{ + Timestamp: e.Timestamp, + Host: host, + Port: port, + Proto: e.Proto, + Method: e.Method, + URL: e.URL, + ReqHeaders: e.ReqHeaders, + ReqBody: e.ReqBody, + StatusCode: e.StatusCode, + RespHeaders: e.RespHeaders, + RespBody: e.RespBody, + DurationMs: e.DurationMs, + TLS: e.TLS, + InScope: s.InScope(host), + }); err != nil { + logger.Error("store: observe response: %v", err) + } +} + +func boolToInt(b bool) int { + if b { + return 1 + } + return 0 +} + +// DeleteTransactionsByHostPort deletes all history rows for the given host and port. +func (s *Store) DeleteTransactionsByHostPort(host string, port int) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM history WHERE host = ? AND port = ?`, host, port) + if err != nil { + return fmt.Errorf("store: delete transactions for %s:%d: %w", host, port, err) + } + return nil + }) +} diff --git a/internal/store/history_test.go b/internal/store/history_test.go new file mode 100644 index 0000000..d3a8e99 --- /dev/null +++ b/internal/store/history_test.go @@ -0,0 +1,190 @@ +package store_test + +import ( + "fmt" + "net/http" + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeTx(method, host, url string, status int) store.Transaction { + return store.Transaction{ + Timestamp: time.Now().UTC(), + Host: host, + Method: method, + URL: url, + ReqHeaders: http.Header{"Content-Type": []string{"text/plain"}}, + ReqBody: []byte("request body"), + StatusCode: status, + RespHeaders: http.Header{"Content-Type": []string{"text/html"}}, + RespBody: []byte("response body"), + DurationMs: 10, + TLS: false, + InScope: true, + } +} + +func TestLog_Insert(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + assert.Equal(t, "GET", txs[0].Method) + assert.Equal(t, "example.com", txs[0].Host) + assert.Equal(t, 200, txs[0].StatusCode) + assert.True(t, txs[0].InScope) +} + +func TestLog_NoDedup_SameMethodHostURLStatus(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + require.NoError(t, st.Log(tx)) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2, "no dedup — every request is a unique row") +} + +func TestLog_NoDedupOnDifferentStatus(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 404))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_NoDedupOnDifferentMethod(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/", 200))) + require.NoError(t, st.Log(makeTx("POST", "example.com", "https://example.com/", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_NoDedupOnDifferentURL(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/a", 200))) + require.NoError(t, st.Log(makeTx("GET", "example.com", "https://example.com/b", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Len(t, txs, 2) +} + +func TestLog_TLSFlagPersisted(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + tx.TLS = true + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + assert.True(t, txs[0].TLS) +} + +func TestAllTransactions_OrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "first.com", "https://first.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "second.com", "https://second.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "third.com", "https://third.com/", 200))) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 3) + assert.Equal(t, "third.com", txs[0].Host) + assert.Equal(t, "second.com", txs[1].Host) + assert.Equal(t, "first.com", txs[2].Host) +} + +func TestAllTransactions_CappedAt500(t *testing.T) { + st := newTestStore(t) + + for i := 0; i < 510; i++ { + url := fmt.Sprintf("https://example.com/path/%d", i) + require.NoError(t, st.Log(makeTx("GET", "example.com", url, 200))) + } + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.LessOrEqual(t, len(txs), 500) +} + +func TestGetTransaction_ReturnsFullBodies(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("POST", "example.com", "https://example.com/api", 201) + tx.ReqBody = []byte(`{"input":"data"}`) + tx.RespBody = []byte(`{"output":"result"}`) + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + + full, err := st.GetTransaction(txs[0].ID) + require.NoError(t, err) + assert.Equal(t, []byte(`{"input":"data"}`), full.ReqBody) + assert.Equal(t, []byte(`{"output":"result"}`), full.RespBody) + assert.Equal(t, "POST", full.Method) + assert.Equal(t, 201, full.StatusCode) +} + +func TestGetTransaction_HeadersDeserialized(t *testing.T) { + st := newTestStore(t) + + tx := makeTx("GET", "example.com", "https://example.com/", 200) + tx.ReqHeaders = http.Header{"X-Request-Id": []string{"req-123"}} + tx.RespHeaders = http.Header{"X-Trace-Id": []string{"trace-456"}} + require.NoError(t, st.Log(tx)) + + txs, err := st.AllTransactions() + require.NoError(t, err) + + full, err := st.GetTransaction(txs[0].ID) + require.NoError(t, err) + assert.Equal(t, "req-123", full.ReqHeaders.Get("X-Request-Id")) + assert.Equal(t, "trace-456", full.RespHeaders.Get("X-Trace-Id")) +} + +func TestGetTransaction_NotFound(t *testing.T) { + st := newTestStore(t) + _, err := st.GetTransaction(99999) + assert.Error(t, err) +} + +func TestClearHistory_RemovesAllRows(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/", 200))) + require.NoError(t, st.Log(makeTx("GET", "b.com", "https://b.com/", 200))) + require.NoError(t, st.ClearHistory()) + + txs, err := st.AllTransactions() + require.NoError(t, err) + assert.Empty(t, txs) +} + +func TestClearHistory_EmptyStoreNoError(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.ClearHistory()) +} diff --git a/internal/store/intercept.go b/internal/store/intercept.go new file mode 100644 index 0000000..e86c683 --- /dev/null +++ b/internal/store/intercept.go @@ -0,0 +1,233 @@ +package store + +import ( + "net" + "net/http" + "sync" + + "github.com/shiv/internal/events" +) + +// PendingRequest is a request held at the intercept gate waiting for a +// user decision. The proxy goroutine blocks on Reply until the UI responds. +// done is closed when the request is cancelled via bypass so watchQueue +// can skip stale entries without deadlocking. +type PendingRequest struct { + Request *http.Request + Body []byte + Reply chan Decision + done chan struct{} +} + +func newPendingRequest(req *http.Request, body []byte) *PendingRequest { + return &PendingRequest{ + Request: req, + Body: body, + Reply: make(chan Decision, 1), + done: make(chan struct{}), + } +} + +// Cancel marks this pending request as cancelled. Safe to call multiple times. +func (p *PendingRequest) Cancel() { + select { + case <-p.done: + // already cancelled + default: + close(p.done) + } +} + +// IsDone returns true if this request was cancelled before the UI processed it. +func (p *PendingRequest) IsDone() bool { + select { + case <-p.done: + return true + default: + return false + } +} + +// Decision is the user's response to a pending request. +type Decision struct { + Forward bool + Request *http.Request + Body []byte +} + +// InterceptGate manages intercept state. When enabled, every request enters +// a serialising semaphore — one request is shown to the UI at a time, others +// wait in order. ForwardAll releases all currently waiting goroutines by +// closing the bypass channel and immediately replacing it so future requests +// are intercepted normally again. +type InterceptGate struct { + mu sync.RWMutex + enabled bool + + queue chan *PendingRequest // UI reads from here + sem chan struct{} // capacity 1 — one request shown at a time + + bypassMu sync.Mutex + bypass chan struct{} // closed to release all waiting goroutines + bypassClosed bool // tracks whether current bypass is already closed + + scopeMu sync.RWMutex + scopeCheckFn func(host string) bool // nil = intercept all; non-nil = intercept only if returns true +} + +// NewInterceptGate creates a new gate. Intercept is off by default. +func NewInterceptGate() *InterceptGate { + return &InterceptGate{ + queue: make(chan *PendingRequest, 1024), + sem: make(chan struct{}, 1), + bypass: make(chan struct{}), + } +} + +// Enabled returns whether intercept is currently on. +func (g *InterceptGate) Enabled() bool { + g.mu.RLock() + defer g.mu.RUnlock() + return g.enabled +} + +// SetEnabled turns intercept on or off. +// Turning off triggers bypass so all goroutines currently blocked in Hold +// are released immediately without deadlocking. +// Turning on resets the bypass channel so future requests are intercepted. +// g.mu and bypassMu are never held simultaneously to avoid deadlock with Hold. +func (g *InterceptGate) SetEnabled(v bool) { + g.mu.Lock() + g.enabled = v + g.mu.Unlock() + + if !v { + g.bypassMu.Lock() + if !g.bypassClosed { + close(g.bypass) + g.bypassClosed = true + } + g.bypassMu.Unlock() + } else { + g.bypassMu.Lock() + g.bypass = make(chan struct{}) + g.bypassClosed = false + g.bypassMu.Unlock() + } +} + +// SetScopeFilter sets a predicate that gates which requests enter the queue. +// Pass nil to intercept all requests. Pass a function returning true for +// in-scope hosts to intercept scope-only traffic. +func (g *InterceptGate) SetScopeFilter(fn func(host string) bool) { + g.scopeMu.Lock() + g.scopeCheckFn = fn + g.scopeMu.Unlock() +} + +// Hold blocks until the user makes a decision on this request. +// If intercept is off, returns immediately. +// Requests acquire the semaphore in arrival order so the UI sees exactly +// one request at a time. If bypass fires at any point during Hold, the +// goroutine exits immediately and the request is passed through unmodified. +// Any pending already in the queue is marked cancelled so watchQueue skips it. +func (g *InterceptGate) Hold(req *http.Request, body []byte) (*http.Request, []byte, bool) { + if !g.Enabled() { + return req, body, true + } + + g.scopeMu.RLock() + fn := g.scopeCheckFn + g.scopeMu.RUnlock() + if fn != nil { + host := req.Host + if h, _, err := net.SplitHostPort(host); err == nil { + host = h + } + if !fn(host) { + return req, body, true + } + } + + // Capture the current bypass reference once. This goroutine watches this + // exact channel throughout — if it gets closed (ForwardAll or SetEnabled(false)), + // we see it no matter where in Hold we are. + g.bypassMu.Lock() + bypass := g.bypass + g.bypassMu.Unlock() + + // Wait for semaphore (one slot) or bypass signal. + select { + case g.sem <- struct{}{}: + case <-bypass: + return req, body, true + } + + // Re-check enabled — may have changed while we waited for the semaphore. + if !g.Enabled() { + <-g.sem + return req, body, true + } + + // Re-check bypass — may have been closed while we raced for the semaphore. + select { + case <-bypass: + <-g.sem + return req, body, true + default: + } + + // We hold the semaphore and bypass is open. Put the request in the queue. + pending := newPendingRequest(req, body) + select { + case g.queue <- pending: + case <-bypass: + pending.Cancel() + <-g.sem + return req, body, true + } + + // Wait for the user's decision or a bypass signal. + select { + case decision := <-pending.Reply: + <-g.sem + return decision.Request, decision.Body, decision.Forward + case <-bypass: + // Bypass fired while we were waiting. Mark the pending as cancelled + // so watchQueue skips this stale queue entry, then release the semaphore. + pending.Cancel() + <-g.sem + return req, body, true + } +} + +// ForwardAll releases all goroutines currently blocked in Hold, passing their +// requests through unmodified. Future requests are intercepted normally — +// the bypass channel is replaced immediately after closing the old one. +func (g *InterceptGate) ForwardAll() { + g.bypassMu.Lock() + if !g.bypassClosed { + close(g.bypass) + g.bypassClosed = true + } + // Replace immediately so future requests are intercepted normally. + g.bypass = make(chan struct{}) + g.bypassClosed = false + g.bypassMu.Unlock() +} + +// Queue returns the channel the UI reads pending requests from. +func (g *InterceptGate) Queue() <-chan *PendingRequest { + return g.queue +} + +// HandleRequest implements events.RequestMiddleware. +// This is a thin adapter: it calls Hold and maps the result to RequestResult. +func (g *InterceptGate) HandleRequest(e events.RequestEvent) events.RequestResult { + req, body, forward := g.Hold(e.Request, e.Body) + return events.RequestResult{ + Drop: !forward, + Request: req, + Body: body, + } +} diff --git a/internal/store/intruder.go b/internal/store/intruder.go new file mode 100644 index 0000000..4968f5d --- /dev/null +++ b/internal/store/intruder.go @@ -0,0 +1,72 @@ +package store + +import "fmt" + +type IntruderConfig struct { + DelayMs int + StopOnStatus int + MaxRedirects int + FollowRedirects string // "never", "always", "in-scope" + TimeoutMs int + RawRequest string + Payloads string +} + +func DefaultIntruderConfig() IntruderConfig { + return IntruderConfig{ + DelayMs: 0, + StopOnStatus: 0, + MaxRedirects: 10, + FollowRedirects: "never", + TimeoutMs: 30000, + RawRequest: "", + Payloads: "", + } +} + +func (s *Store) LoadIntruderConfig() IntruderConfig { + config := DefaultIntruderConfig() + err := s.db.QueryRow(` + SELECT delay_ms, stop_on_status, max_redirects, follow_redirects, timeout_ms, raw_request, payloads + FROM intruder_config WHERE id = 1`).Scan( + &config.DelayMs, + &config.StopOnStatus, + &config.MaxRedirects, + &config.FollowRedirects, + &config.TimeoutMs, + &config.RawRequest, + &config.Payloads, + ) + if err != nil { + return config + } + return config +} + +func (s *Store) SaveIntruderConfig(config IntruderConfig) error { + return s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO intruder_config (id, delay_ms, stop_on_status, max_redirects, follow_redirects, timeout_ms, raw_request, payloads) + VALUES (1, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(id) DO UPDATE SET + delay_ms = excluded.delay_ms, + stop_on_status = excluded.stop_on_status, + max_redirects = excluded.max_redirects, + follow_redirects = excluded.follow_redirects, + timeout_ms = excluded.timeout_ms, + raw_request = excluded.raw_request, + payloads = excluded.payloads`, + config.DelayMs, + config.StopOnStatus, + config.MaxRedirects, + config.FollowRedirects, + config.TimeoutMs, + config.RawRequest, + config.Payloads, + ) + if err != nil { + return fmt.Errorf("store: save intruder config: %w", err) + } + return nil + }) +} diff --git a/internal/store/loot.go b/internal/store/loot.go new file mode 100644 index 0000000..c610db1 --- /dev/null +++ b/internal/store/loot.go @@ -0,0 +1,103 @@ +package store + +import ( + "fmt" + "time" +) + +type LootEntry struct { + ID int64 + Title string + Severity string + Notes string + HistoryID *uint64 + RawRequest string + RawResponse string + CreatedAt time.Time +} + +func (s *Store) AllLoot() ([]LootEntry, error) { + rows, err := s.db.Query(` + SELECT id, title, severity, notes, history_id, raw_request, raw_response, created_at + FROM loot ORDER BY + CASE severity + WHEN 'Critical' THEN 1 + WHEN 'High' THEN 2 + WHEN 'Medium' THEN 3 + WHEN 'Low' THEN 4 + WHEN 'Info' THEN 5 + ELSE 6 + END ASC, id DESC`) + if err != nil { + return nil, fmt.Errorf("store: query loot: %w", err) + } + defer rows.Close() + + var entries []LootEntry + for rows.Next() { + var entry LootEntry + var histID *int64 + var timestampStr string + if err := rows.Scan(&entry.ID, &entry.Title, &entry.Severity, &entry.Notes, &histID, &entry.RawRequest, &entry.RawResponse, ×tampStr); err != nil { + return nil, fmt.Errorf("store: scan loot: %w", err) + } + if histID != nil { + histIDUint := uint64(*histID) + entry.HistoryID = &histIDUint + } + entry.CreatedAt, _ = time.Parse(time.RFC3339, timestampStr) + entries = append(entries, entry) + } + return entries, rows.Err() +} + +func (s *Store) AddLoot(entry LootEntry) (int64, error) { + var id int64 + err := s.write(func() error { + var histID any + if entry.HistoryID != nil { + histID = *entry.HistoryID + } + result, err := s.db.Exec(` + INSERT INTO loot (title, severity, notes, history_id, raw_request, raw_response, created_at) + VALUES (?, ?, ?, ?, ?, ?, ?)`, + entry.Title, entry.Severity, entry.Notes, histID, + entry.RawRequest, entry.RawResponse, + time.Now().UTC().Format(time.RFC3339), + ) + if err != nil { + return fmt.Errorf("store: add loot: %w", err) + } + id, err = result.LastInsertId() + return err + }) + if err != nil { + return 0, err + } + s.notifyLoot() + return id, nil +} + +func (s *Store) DeleteLoot(id int64) error { + err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM loot WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete loot: %w", err) + } + return nil + }) + if err != nil { + return err + } + s.notifyLoot() + return nil +} + +// notifyLoot signals the UI that loot state has changed. +// Non-blocking — a full buffer means the UI is already scheduled to reload. +func (s *Store) notifyLoot() { + select { + case s.LootEntries <- struct{}{}: + default: + } +} diff --git a/internal/store/loot_test.go b/internal/store/loot_test.go new file mode 100644 index 0000000..59bd4fe --- /dev/null +++ b/internal/store/loot_test.go @@ -0,0 +1,155 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeLoot(title, severity, notes string) store.LootEntry { + return store.LootEntry{ + Title: title, + Severity: severity, + Notes: notes, + } +} + +func TestAddLoot_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + entry := store.LootEntry{ + Title: "Admin panel exposed", + Severity: "High", + Notes: "Found at /admin with no auth", + } + id, err := st.AddLoot(entry) + require.NoError(t, err) + assert.Positive(t, id) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + + got := entries[0] + assert.Equal(t, id, got.ID) + assert.Equal(t, "Admin panel exposed", got.Title) + assert.Equal(t, "High", got.Severity) + assert.Equal(t, "Found at /admin with no auth", got.Notes) + assert.Nil(t, got.HistoryID) + assert.WithinDuration(t, time.Now(), got.CreatedAt, 5*time.Second) +} + +func TestAddLoot_WithLinkedHistoryID(t *testing.T) { + st := newTestStore(t) + + // first log a transaction to get a valid history ID + tx := makeTx("GET", "example.com", "https://example.com/secret", 200) + require.NoError(t, st.Log(tx)) + txs, err := st.AllTransactions() + require.NoError(t, err) + require.Len(t, txs, 1) + histID := txs[0].ID + + entry := store.LootEntry{ + Title: "Secret endpoint", + Severity: "Critical", + Notes: "Returns sensitive data", + HistoryID: &histID, + } + id, err := st.AddLoot(entry) + require.NoError(t, err) + assert.Positive(t, id) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + require.NotNil(t, entries[0].HistoryID) + assert.Equal(t, histID, *entries[0].HistoryID) +} + +func TestAllLoot_OrderedBySeverityThenIDDesc(t *testing.T) { + st := newTestStore(t) + + // Insert in reverse severity order to confirm sorting + _, err := st.AddLoot(makeLoot("Info finding", "Info", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Low finding", "Low", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Critical finding", "Critical", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("High finding", "High", "")) + require.NoError(t, err) + _, err = st.AddLoot(makeLoot("Medium finding", "Medium", "")) + require.NoError(t, err) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 5) + + assert.Equal(t, "Critical", entries[0].Severity) + assert.Equal(t, "High", entries[1].Severity) + assert.Equal(t, "Medium", entries[2].Severity) + assert.Equal(t, "Low", entries[3].Severity) + assert.Equal(t, "Info", entries[4].Severity) +} + +func TestAllLoot_MultipleSameSeverityOrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + id1, err := st.AddLoot(makeLoot("First High", "High", "")) + require.NoError(t, err) + id2, err := st.AddLoot(makeLoot("Second High", "High", "")) + require.NoError(t, err) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 2) + + // same severity — ordered by id DESC so id2 comes first + assert.Equal(t, id2, entries[0].ID) + assert.Equal(t, id1, entries[1].ID) +} + +func TestDeleteLoot(t *testing.T) { + st := newTestStore(t) + + id, err := st.AddLoot(makeLoot("To delete", "Low", "")) + require.NoError(t, err) + + require.NoError(t, st.DeleteLoot(id)) + + entries, err := st.AllLoot() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestDeleteLoot_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteLoot(9999)) +} + +func TestDeleteLoot_DoesNotAffectOtherEntries(t *testing.T) { + st := newTestStore(t) + + idKeep, err := st.AddLoot(makeLoot("Keep", "High", "")) + require.NoError(t, err) + idDel, err := st.AddLoot(makeLoot("Delete", "Low", "")) + require.NoError(t, err) + + require.NoError(t, st.DeleteLoot(idDel)) + + entries, err := st.AllLoot() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, idKeep, entries[0].ID) +} + +func TestAllLoot_Empty(t *testing.T) { + st := newTestStore(t) + entries, err := st.AllLoot() + require.NoError(t, err) + assert.Empty(t, entries) +} diff --git a/internal/store/plugin.go b/internal/store/plugin.go new file mode 100644 index 0000000..545d9b4 --- /dev/null +++ b/internal/store/plugin.go @@ -0,0 +1,156 @@ +package store + +import ( + "fmt" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +const maxPluginLogLines = 500 + +// PluginEntry represents a loaded plugin and its persistent enabled state. +type PluginEntry struct { + Name string + Path string + Enabled bool +} + +// RegisterPlugin records a plugin in the database. If the plugin is new, +// enabled defaults to true. If it already exists, only the path is updated — +// the existing enabled state is preserved. Notifies the UI via PluginEntries. +func (s *Store) RegisterPlugin(name, path string) error { + if err := s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO plugins (name, path, enabled) VALUES (?, ?, 1) + ON CONFLICT(name) DO UPDATE SET path = excluded.path`, + name, path, + ) + if err != nil { + return fmt.Errorf("store: register plugin %s: %w", name, err) + } + return nil + }); err != nil { + return err + } + + select { + case s.PluginEntries <- struct{}{}: + default: + } + return nil +} + +// Scans for plugins on the disk, removes them from DB if they are deleted from disk. +func (s *Store) ScanPlugins(active []string) error { + set := make(map[string]struct{}, len(active)) + for _, name := range active { + set[name] = struct{}{} + } + all, err := s.AllPlugins() + if err != nil { + return err + } + for _, e := range all { + if _, ok := set[e.Name]; ok { + continue + } + if err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM plugins WHERE name = ?`, e.Name) + return err + }); err != nil { + return fmt.Errorf("store: scan plugins remove %s: %w", e.Name, err) + } + } + select { + case s.PluginEntries <- struct{}{}: + default: + } + return nil +} + +// PluginEnabled returns the persisted enabled state for a plugin. +// Returns true if the plugin record is not found (safe default for new plugins). +func (s *Store) PluginEnabled(name string) (bool, error) { + var enabled int + err := s.db.QueryRow(`SELECT enabled FROM plugins WHERE name = ?`, name).Scan(&enabled) + if err != nil { + return true, err + } + return enabled == 1, nil +} + +// AllPlugins returns all known plugins ordered by name. +func (s *Store) AllPlugins() ([]PluginEntry, error) { + rows, err := s.db.Query(`SELECT name, path, enabled FROM plugins ORDER BY name ASC`) + if err != nil { + return nil, fmt.Errorf("store: query plugins: %w", err) + } + defer rows.Close() + + var entries []PluginEntry + for rows.Next() { + var e PluginEntry + var enabledInt int + if err := rows.Scan(&e.Name, &e.Path, &enabledInt); err != nil { + return nil, fmt.Errorf("store: scan plugin: %w", err) + } + e.Enabled = enabledInt == 1 + entries = append(entries, e) + } + return entries, rows.Err() +} + +// PluginLogs returns the in-memory log lines for a plugin, oldest first. +// Returns a copy to avoid races with the writer. +func (s *Store) PluginLogs(name string) []string { + s.pluginLogsMu.RLock() + defer s.pluginLogsMu.RUnlock() + lines := s.pluginLogs[name] + if len(lines) == 0 { + return nil + } + result := make([]string, len(lines)) + copy(result, lines) + return result +} + +// ObservePluginLog implements events.PluginLogObserver. +// Appends the log line to the in-memory ring buffer and notifies the UI. +func (s *Store) ObservePluginLog(e events.PluginLogEvent) { + s.pluginLogsMu.Lock() + lines := s.pluginLogs[e.Name] + lines = append(lines, e.Message) + if len(lines) > maxPluginLogLines { + lines = lines[len(lines)-maxPluginLogLines:] + } + s.pluginLogs[e.Name] = lines + s.pluginLogsMu.Unlock() + + select { + case s.PluginEntries <- struct{}{}: + default: + } +} + +// ObservePluginEnabled implements events.PluginEnabledObserver. +// Persists the new enabled state to the database and notifies the UI. +func (s *Store) ObservePluginEnabled(e events.SetPluginEnabledEvent) { + if err := s.write(func() error { + _, err := s.db.Exec(`UPDATE plugins SET enabled = ? WHERE name = ?`, + boolToInt(e.Enabled), e.Name, + ) + if err != nil { + return fmt.Errorf("store: set plugin enabled %s: %w", e.Name, err) + } + return nil + }); err != nil { + logger.Error("store: observe plugin enabled: %v", err) + return + } + + select { + case s.PluginEntries <- struct{}{}: + default: + } +} diff --git a/internal/store/plugin_test.go b/internal/store/plugin_test.go new file mode 100644 index 0000000..b640628 --- /dev/null +++ b/internal/store/plugin_test.go @@ -0,0 +1,368 @@ +package store_test + +import ( + "fmt" + "sync" + "testing" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// ── RegisterPlugin ──────────────────────────────────────────────────────────── + +func TestRegisterPlugin_NewPlugin_DefaultsToEnabled(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "test.lua", entries[0].Name) + assert.Equal(t, "/plugins/test.lua", entries[0].Path) + assert.True(t, entries[0].Enabled) +} + +func TestRegisterPlugin_ExistingPlugin_PreservesEnabledState(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + // Re-register with a new path — enabled state must be preserved. + require.NoError(t, st.RegisterPlugin("test.lua", "/new/path/test.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "/new/path/test.lua", entries[0].Path) + assert.False(t, entries[0].Enabled, "enabled state must survive re-registration") +} + +func TestRegisterPlugin_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + + // Drain any startup notifications. + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after RegisterPlugin") + } +} + +func TestRegisterPlugin_MultiplePlugins(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.RegisterPlugin("a.lua", "/plugins/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/plugins/b.lua")) + require.NoError(t, st.RegisterPlugin("c.lua", "/plugins/c.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Len(t, entries, 3) +} + +// ── PluginEnabled ───────────────────────────────────────────────────────────── + +func TestPluginEnabled_UnknownPlugin_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + + enabled, err := st.PluginEnabled("nonexistent.lua") + // Error is expected (no row) but the safe default is true. + assert.True(t, enabled) + _ = err +} + +func TestPluginEnabled_NewPlugin_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +func TestPluginEnabled_AfterDisable_ReturnsFalse(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.False(t, enabled) +} + +func TestPluginEnabled_AfterReEnable_ReturnsTrue(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/plugins/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +// ── AllPlugins ──────────────────────────────────────────────────────────────── + +func TestAllPlugins_Empty(t *testing.T) { + st := newTestStore(t) + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestAllPlugins_OrderedByNameAsc(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("zebra.lua", "/z.lua")) + require.NoError(t, st.RegisterPlugin("alpha.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("mango.lua", "/m.lua")) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 3) + assert.Equal(t, "alpha.lua", entries[0].Name) + assert.Equal(t, "mango.lua", entries[1].Name) + assert.Equal(t, "zebra.lua", entries[2].Name) +} + +func TestAllPlugins_EnabledFieldReflectsState(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("on.lua", "/on.lua")) + require.NoError(t, st.RegisterPlugin("off.lua", "/off.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "off.lua", Enabled: false}) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 2) + + byName := make(map[string]store.PluginEntry) + for _, e := range entries { + byName[e.Name] = e + } + assert.True(t, byName["on.lua"].Enabled) + assert.False(t, byName["off.lua"].Enabled) +} + +// ── ScanPlugins ─────────────────────────────────────────────────────────────── + +func TestScanPlugins_RemovesStaleRecords(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("keep.lua", "/keep.lua")) + require.NoError(t, st.RegisterPlugin("gone.lua", "/gone.lua")) + + require.NoError(t, st.ScanPlugins([]string{"keep.lua"})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "keep.lua", entries[0].Name) +} + +func TestScanPlugins_EmptyActive_RemovesAll(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("a.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/b.lua")) + + require.NoError(t, st.ScanPlugins([]string{})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestScanPlugins_AllPresent_NothingRemoved(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("a.lua", "/a.lua")) + require.NoError(t, st.RegisterPlugin("b.lua", "/b.lua")) + + require.NoError(t, st.ScanPlugins([]string{"a.lua", "b.lua"})) + + entries, err := st.AllPlugins() + require.NoError(t, err) + assert.Len(t, entries, 2) +} + +func TestScanPlugins_EmptyStore_NoPanic(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.ScanPlugins([]string{"a.lua"})) +} + +func TestScanPlugins_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("gone.lua", "/gone.lua")) + + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + require.NoError(t, st.ScanPlugins([]string{})) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ScanPlugins") + } +} + +// ── PluginLogs ──────────────────────────────────────────────────────────────── + +func TestPluginLogs_UnknownPlugin_ReturnsNil(t *testing.T) { + st := newTestStore(t) + assert.Nil(t, st.PluginLogs("nonexistent.lua")) +} + +func TestPluginLogs_ReturnsLinesInOrder(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 1"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 2"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "line 3"}) + + lines := st.PluginLogs("test.lua") + require.Len(t, lines, 3) + assert.Equal(t, "line 1", lines[0]) + assert.Equal(t, "line 2", lines[1]) + assert.Equal(t, "line 3", lines[2]) +} + +func TestPluginLogs_ReturnsACopy(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "original"}) + + lines := st.PluginLogs("test.lua") + lines[0] = "mutated" + + // Original must be unchanged. + assert.Equal(t, "original", st.PluginLogs("test.lua")[0]) +} + +func TestPluginLogs_IsolatedPerPlugin(t *testing.T) { + st := newTestStore(t) + st.ObservePluginLog(events.PluginLogEvent{Name: "a.lua", Message: "from a"}) + st.ObservePluginLog(events.PluginLogEvent{Name: "b.lua", Message: "from b"}) + + assert.Equal(t, []string{"from a"}, st.PluginLogs("a.lua")) + assert.Equal(t, []string{"from b"}, st.PluginLogs("b.lua")) +} + +func TestPluginLogs_RingBuffer_CapsAtMaxLines(t *testing.T) { + st := newTestStore(t) + // Write more than maxPluginLogLines (500) entries. + for i := 0; i < 510; i++ { + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: fmt.Sprintf("line %d", i)}) + } + + lines := st.PluginLogs("test.lua") + assert.Len(t, lines, 500) + // The oldest 10 lines (0-9) were evicted; buffer starts at line 10. + assert.Equal(t, "line 10", lines[0]) + assert.Equal(t, "line 509", lines[499]) +} + +// ── ObservePluginLog ────────────────────────────────────────────────────────── + +func TestObservePluginLog_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "hello"}) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ObservePluginLog") + } +} + +func TestObservePluginLog_ChannelNonBlocking_WhenFull(t *testing.T) { + st := newTestStore(t) + // Fill the channel to capacity. + for i := 0; i < cap(st.PluginEntries); i++ { + st.PluginEntries <- struct{}{} + } + // Must not block even when channel is full. + assert.NotPanics(t, func() { + st.ObservePluginLog(events.PluginLogEvent{Name: "test.lua", Message: "msg"}) + }) +} + +func TestObservePluginLog_ConcurrentWrites_Safe(t *testing.T) { + st := newTestStore(t) + const goroutines = 50 + var wg sync.WaitGroup + wg.Add(goroutines) + for i := 0; i < goroutines; i++ { + go func(i int) { + defer wg.Done() + st.ObservePluginLog(events.PluginLogEvent{ + Name: "test.lua", + Message: fmt.Sprintf("msg %d", i), + }) + }(i) + } + wg.Wait() + + lines := st.PluginLogs("test.lua") + assert.Len(t, lines, goroutines) +} + +// ── ObservePluginEnabled ────────────────────────────────────────────────────── + +func TestObservePluginEnabled_PersistsDisabled(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.False(t, enabled) +} + +func TestObservePluginEnabled_PersistsEnabled(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: true}) + + enabled, err := st.PluginEnabled("test.lua") + require.NoError(t, err) + assert.True(t, enabled) +} + +func TestObservePluginEnabled_NotifiesPluginEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.RegisterPlugin("test.lua", "/test.lua")) + + for len(st.PluginEntries) > 0 { + <-st.PluginEntries + } + + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "test.lua", Enabled: false}) + + select { + case <-st.PluginEntries: + default: + t.Fatal("PluginEntries channel was not notified after ObservePluginEnabled") + } +} + +func TestObservePluginEnabled_UnknownPlugin_NoError(t *testing.T) { + st := newTestStore(t) + // Toggling a plugin not in the DB must not panic or crash. + assert.NotPanics(t, func() { + st.ObservePluginEnabled(events.SetPluginEnabledEvent{Name: "ghost.lua", Enabled: false}) + }) +} diff --git a/internal/store/project.go b/internal/store/project.go new file mode 100644 index 0000000..e379f32 --- /dev/null +++ b/internal/store/project.go @@ -0,0 +1,245 @@ +package store + +import ( + "database/sql" + "fmt" + "sync" + + _ "modernc.org/sqlite" +) + +type Store struct { + db *sql.DB + writeCh chan func() error + done chan struct{} + + Updates chan Transaction + Intercept *InterceptGate + WebSocketConnections chan WebSocketConnection + WebSocketFrames chan WebSocketFrame + LootEntries chan struct{} + PluginEntries chan struct{} + + // pluginLogs holds in-memory log lines per plugin name. + // Protected by pluginLogsMu — separate from the SQLite write serialiser + // because log appends must never block on DB I/O. + pluginLogsMu sync.RWMutex + pluginLogs map[string][]string + + // scopeCache is an in-memory copy of the scope table. + // Populated lazily on first InScope call and invalidated on every write. + // Protected by scopeCacheMu — separate from the write serialiser so + // InScope reads never block on DB I/O. + scopeCacheMu sync.RWMutex + scopeCache []ScopeEntry + scopeCacheSet bool +} + +func Open(path string) (*Store, error) { + db, err := sql.Open("sqlite", path) + if err != nil { + return nil, fmt.Errorf("store: open %s: %w", path, err) + } + db.SetMaxOpenConns(1) + if _, err := db.Exec(`PRAGMA journal_mode=WAL`); err != nil { + return nil, fmt.Errorf("store: enable WAL: %w", err) + } + + projectStore := &Store{ + db: db, + writeCh: make(chan func() error, 256), + done: make(chan struct{}), + Updates: make(chan Transaction, 256), + Intercept: NewInterceptGate(), + WebSocketConnections: make(chan WebSocketConnection, 256), + WebSocketFrames: make(chan WebSocketFrame, 256), + LootEntries: make(chan struct{}, 256), + PluginEntries: make(chan struct{}, 256), + pluginLogs: make(map[string][]string), + } + if err := projectStore.migrate(); err != nil { + db.Close() + return nil, err + } + + go projectStore.writeLoop() + return projectStore, nil +} + +func (s *Store) Close() error { + close(s.done) + close(s.Updates) + close(s.WebSocketConnections) + close(s.WebSocketFrames) + close(s.LootEntries) + close(s.PluginEntries) + close(s.Intercept.queue) + return s.db.Close() +} + +func (s *Store) write(writeFunc func() error) error { + errorChan := make(chan error, 1) + select { + case s.writeCh <- func() error { + err := writeFunc() + errorChan <- err + return err + }: + case <-s.done: + return fmt.Errorf("store: closed") + } + select { + case err := <-errorChan: + return err + case <-s.done: + return fmt.Errorf("store: closed") + } +} + +func (s *Store) writeLoop() { + for { + select { + case fn := <-s.writeCh: + _ = fn() + case <-s.done: + + for { + select { + case fn := <-s.writeCh: + _ = fn() + default: + return + } + } + } + } +} + +func (s *Store) migrate() error { + _, err := s.db.Exec(schema) + if err != nil { + return fmt.Errorf("store: migrate: %w", err) + } + // Add new columns to existing tables if they don't exist. + // Errors are ignored — column may already exist. + migrations := []string{ + `ALTER TABLE loot ADD COLUMN raw_request TEXT DEFAULT ''`, + `ALTER TABLE loot ADD COLUMN raw_response TEXT DEFAULT ''`, + `ALTER TABLE history ADD COLUMN proto TEXT NOT NULL DEFAULT 'HTTP/1.1'`, + `ALTER TABLE history ADD COLUMN port INTEGER NOT NULL DEFAULT 443`, + `ALTER TABLE repeater_tabs ADD COLUMN tab_type TEXT NOT NULL DEFAULT 'http'`, + `ALTER TABLE history ADD COLUMN resp_size INTEGER NOT NULL DEFAULT 0`, + `CREATE INDEX IF NOT EXISTS idx_history_host_port_id ON history(host, port, id DESC)`, + } + for _, migration := range migrations { + s.db.Exec(migration) + } + return nil +} + +const schema = ` +CREATE TABLE IF NOT EXISTS meta ( + key TEXT PRIMARY KEY, + value TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS history ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + timestamp TEXT NOT NULL, + host TEXT NOT NULL, + port INTEGER NOT NULL DEFAULT 443, + method TEXT NOT NULL, + url TEXT NOT NULL, + proto TEXT NOT NULL DEFAULT 'HTTP/1.1', + req_headers TEXT NOT NULL, + req_body BLOB, + status_code INTEGER, + resp_headers TEXT, + resp_body BLOB, + resp_size INTEGER NOT NULL DEFAULT 0, + duration_ms INTEGER, + tls INTEGER DEFAULT 0, + in_scope INTEGER DEFAULT 1 +); + +CREATE INDEX IF NOT EXISTS idx_history_host ON history(host); +CREATE INDEX IF NOT EXISTS idx_history_timestamp ON history(timestamp); +CREATE INDEX IF NOT EXISTS idx_history_scope ON history(in_scope); +CREATE INDEX IF NOT EXISTS idx_history_dedup ON history(method, host, url, status_code); +CREATE INDEX IF NOT EXISTS idx_history_host_id ON history(host, id DESC); +CREATE INDEX IF NOT EXISTS idx_history_scope_id ON history(in_scope, id DESC); + +CREATE TABLE IF NOT EXISTS repeater_tabs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL DEFAULT 'Tab', + host TEXT NOT NULL, + port INTEGER NOT NULL DEFAULT 80, + tls INTEGER DEFAULT 0, + raw_request TEXT NOT NULL, + last_response TEXT, + position INTEGER DEFAULT 0 +); + +CREATE TABLE IF NOT EXISTS loot ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + title TEXT NOT NULL, + severity TEXT NOT NULL, + notes TEXT, + history_id INTEGER REFERENCES history(id), + raw_request TEXT DEFAULT '', + raw_response TEXT DEFAULT '', + created_at TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS scope ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + host TEXT NOT NULL UNIQUE +); + +CREATE TABLE IF NOT EXISTS intruder_config ( + id INTEGER PRIMARY KEY, + delay_ms INTEGER DEFAULT 0, + stop_on_status INTEGER DEFAULT 0, + max_redirects INTEGER DEFAULT 10, + follow_redirects TEXT DEFAULT 'never', + timeout_ms INTEGER DEFAULT 30000, + raw_request TEXT DEFAULT '', + payloads TEXT DEFAULT '' +); + +CREATE TABLE IF NOT EXISTS websocket_connections ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + history_id INTEGER REFERENCES history(id), + host TEXT NOT NULL, + url TEXT NOT NULL, + tls INTEGER DEFAULT 0, + in_scope INTEGER DEFAULT 1, + timestamp TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_ws_connections_host ON websocket_connections(host); +CREATE INDEX IF NOT EXISTS idx_ws_connections_scope ON websocket_connections(in_scope); + +CREATE TABLE IF NOT EXISTS websocket_frames ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + connection_id INTEGER NOT NULL REFERENCES websocket_connections(id) ON DELETE CASCADE, + timestamp TEXT NOT NULL, + direction INTEGER NOT NULL, -- 0 = client→server, 1 = server→client + opcode INTEGER NOT NULL, + payload BLOB +); + +CREATE INDEX IF NOT EXISTS idx_ws_frames_connection ON websocket_frames(connection_id); + +CREATE TABLE IF NOT EXISTS plugins ( + name TEXT PRIMARY KEY, + path TEXT NOT NULL DEFAULT '', + enabled INTEGER NOT NULL DEFAULT 1 +); + +CREATE TABLE IF NOT EXISTS annotations ( + history_id INTEGER PRIMARY KEY REFERENCES history(id) ON DELETE CASCADE, + comment TEXT NOT NULL DEFAULT '', + color TEXT NOT NULL DEFAULT '' +); +` diff --git a/internal/store/repeater.go b/internal/store/repeater.go new file mode 100644 index 0000000..b433559 --- /dev/null +++ b/internal/store/repeater.go @@ -0,0 +1,103 @@ +package store + +import ( + "fmt" +) + +// RepeaterTab represents a saved repeater tab. +type RepeaterTab struct { + ID int64 + Name string + Host string + Port int + TLS bool + TabType string // "http" or "websocket" + RawRequest string + LastResponse string + Position int +} + +// AllRepeaterTabs returns all saved repeater tabs ordered by position. +func (s *Store) AllRepeaterTabs() ([]RepeaterTab, error) { + rows, err := s.db.Query(` + SELECT id, name, host, port, tls, COALESCE(tab_type, 'http'), raw_request, COALESCE(last_response, ''), position + FROM repeater_tabs + ORDER BY position ASC`) + if err != nil { + return nil, fmt.Errorf("store: query repeater tabs: %w", err) + } + defer rows.Close() + + var tabs []RepeaterTab + for rows.Next() { + var tab RepeaterTab + var tlsFlag int + if err := rows.Scan(&tab.ID, &tab.Name, &tab.Host, &tab.Port, &tlsFlag, &tab.TabType, &tab.RawRequest, &tab.LastResponse, &tab.Position); err != nil { + return nil, fmt.Errorf("store: scan repeater tab: %w", err) + } + tab.TLS = tlsFlag == 1 + if tab.TabType == "" { + tab.TabType = "http" + } + tabs = append(tabs, tab) + } + return tabs, rows.Err() +} + +// SaveRepeaterTab inserts a new repeater tab and returns its ID. +func (s *Store) SaveRepeaterTab(tab RepeaterTab) (int64, error) { + tabType := tab.TabType + if tabType == "" { + tabType = "http" + } + var id int64 + err := s.write(func() error { + result, err := s.db.Exec(` + INSERT INTO repeater_tabs (name, host, port, tls, tab_type, raw_request, last_response, position) + VALUES (?, ?, ?, ?, ?, ?, ?, (SELECT COALESCE(MAX(position), 0) + 1 FROM repeater_tabs))`, + tab.Name, tab.Host, tab.Port, boolToInt(tab.TLS), tabType, tab.RawRequest, tab.LastResponse, + ) + if err != nil { + return fmt.Errorf("store: save repeater tab: %w", err) + } + id, err = result.LastInsertId() + return err + }) + return id, err +} + +// UpdateRepeaterTab updates the request and response for an existing tab. +func (s *Store) UpdateRepeaterTab(id int64, rawRequest, lastResponse string) error { + return s.write(func() error { + _, err := s.db.Exec(` + UPDATE repeater_tabs SET raw_request = ?, last_response = ? WHERE id = ?`, + rawRequest, lastResponse, id, + ) + if err != nil { + return fmt.Errorf("store: update repeater tab: %w", err) + } + return nil + }) +} + +// RenameRepeaterTab updates the name of a tab. +func (s *Store) RenameRepeaterTab(id int64, name string) error { + return s.write(func() error { + _, err := s.db.Exec(`UPDATE repeater_tabs SET name = ? WHERE id = ?`, name, id) + if err != nil { + return fmt.Errorf("store: rename repeater tab: %w", err) + } + return nil + }) +} + +// DeleteRepeaterTab removes a repeater tab by ID. +func (s *Store) DeleteRepeaterTab(id int64) error { + return s.write(func() error { + _, err := s.db.Exec(`DELETE FROM repeater_tabs WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete repeater tab: %w", err) + } + return nil + }) +} diff --git a/internal/store/repeater_test.go b/internal/store/repeater_test.go new file mode 100644 index 0000000..c85ad83 --- /dev/null +++ b/internal/store/repeater_test.go @@ -0,0 +1,265 @@ +package store_test + +import ( + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func makeTab(name, host string, port int, tls bool) store.RepeaterTab { + return store.RepeaterTab{ + Name: name, + Host: host, + Port: port, + TLS: tls, + TabType: "http", + RawRequest: "GET / HTTP/1.1\r\nHost: " + host + "\r\n\r\n", + } +} + +func makeWSTab(name, url string, tls bool) store.RepeaterTab { + return store.RepeaterTab{ + Name: name, + Host: url, + Port: 0, + TLS: tls, + TabType: "websocket", + RawRequest: url, + } +} + +// ── HTTP tab ────────────────────────────────────────────────────────────────── + +func TestSaveRepeaterTab_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + tab := store.RepeaterTab{ + Name: "My Tab", + Host: "example.com", + Port: 443, + TLS: true, + TabType: "http", + RawRequest: "GET /api HTTP/1.1\r\nHost: example.com\r\n\r\n", + } + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + + got := tabs[0] + assert.Equal(t, id, got.ID) + assert.Equal(t, "My Tab", got.Name) + assert.Equal(t, "example.com", got.Host) + assert.Equal(t, 443, got.Port) + assert.True(t, got.TLS) + assert.Equal(t, "http", got.TabType) + assert.Equal(t, tab.RawRequest, got.RawRequest) +} + +func TestSaveRepeaterTab_TLSFalse(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("HTTP Tab", "example.com", 80, false)) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.False(t, tabs[0].TLS) + assert.Equal(t, "http", tabs[0].TabType) +} + +func TestSaveRepeaterTab_DefaultTabType(t *testing.T) { + st := newTestStore(t) + + // TabType omitted — should default to "http". + tab := store.RepeaterTab{ + Name: "No Type", + Host: "example.com", + Port: 80, + RawRequest: "GET / HTTP/1.1\r\nHost: example.com\r\n\r\n", + } + _, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "http", tabs[0].TabType) +} + +func TestAllRepeaterTabs_OrderedByPosition(t *testing.T) { + st := newTestStore(t) + + for _, name := range []string{"First", "Second", "Third"} { + _, err := st.SaveRepeaterTab(makeTab(name, "example.com", 80, false)) + require.NoError(t, err) + } + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 3) + assert.Equal(t, "First", tabs[0].Name) + assert.Equal(t, "Second", tabs[1].Name) + assert.Equal(t, "Third", tabs[2].Name) +} + +func TestAllRepeaterTabs_PositionAutoIncrements(t *testing.T) { + st := newTestStore(t) + + _, err := st.SaveRepeaterTab(makeTab("A", "a.com", 80, false)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeTab("B", "b.com", 80, false)) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 2) + assert.Less(t, tabs[0].Position, tabs[1].Position) +} + +func TestUpdateRepeaterTab_UpdatesRequestAndResponse(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Tab", "example.com", 80, false)) + require.NoError(t, err) + + newReq := "POST /submit HTTP/1.1\r\nHost: example.com\r\n\r\nbody" + newResp := "HTTP/1.1 200 OK\r\n\r\nok" + require.NoError(t, st.UpdateRepeaterTab(id, newReq, newResp)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, newReq, tabs[0].RawRequest) + assert.Equal(t, newResp, tabs[0].LastResponse) +} + +func TestUpdateRepeaterTab_NonExistentIDNoError(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.UpdateRepeaterTab(9999, "req", "resp")) +} + +func TestRenameRepeaterTab(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Old", "example.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.RenameRepeaterTab(id, "New Name")) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "New Name", tabs[0].Name) +} + +func TestDeleteRepeaterTab(t *testing.T) { + st := newTestStore(t) + + id, err := st.SaveRepeaterTab(makeTab("Tab", "example.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(id)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + assert.Empty(t, tabs) +} + +func TestDeleteRepeaterTab_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteRepeaterTab(9999)) +} + +func TestDeleteRepeaterTab_DoesNotAffectOtherTabs(t *testing.T) { + st := newTestStore(t) + + idKeep, err := st.SaveRepeaterTab(makeTab("Keep", "keep.com", 80, false)) + require.NoError(t, err) + idDel, err := st.SaveRepeaterTab(makeTab("Delete", "del.com", 80, false)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(idDel)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, idKeep, tabs[0].ID) + assert.Equal(t, "Keep", tabs[0].Name) +} + +// ── WebSocket tab ───────────────────────────────────────────────────────────── + +func TestSaveRepeaterTab_WebSocketType(t *testing.T) { + st := newTestStore(t) + + tab := makeWSTab("WS Tab", "wss://echo.websocket.org:443/", true) + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + assert.Positive(t, id) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, "websocket", tabs[0].TabType) + assert.Equal(t, "wss://echo.websocket.org:443/", tabs[0].RawRequest) + assert.True(t, tabs[0].TLS) +} + +func TestAllRepeaterTabs_MixedTypes(t *testing.T) { + st := newTestStore(t) + + _, err := st.SaveRepeaterTab(makeTab("HTTP", "example.com", 80, false)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeWSTab("WS", "wss://example.com/ws", true)) + require.NoError(t, err) + _, err = st.SaveRepeaterTab(makeTab("HTTP2", "other.com", 443, true)) + require.NoError(t, err) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 3) + assert.Equal(t, "http", tabs[0].TabType) + assert.Equal(t, "websocket", tabs[1].TabType) + assert.Equal(t, "http", tabs[2].TabType) +} + +func TestSaveRepeaterTab_WebSocketPreservesURL(t *testing.T) { + st := newTestStore(t) + + url := "wss://api.example.com:8443/realtime?token=abc123" + tab := makeWSTab("API WS", url, true) + id, err := st.SaveRepeaterTab(tab) + require.NoError(t, err) + + full, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, full, 1) + assert.Equal(t, id, full[0].ID) + assert.Equal(t, url, full[0].RawRequest) + assert.Equal(t, "websocket", full[0].TabType) +} + +func TestDeleteRepeaterTab_WebSocket(t *testing.T) { + st := newTestStore(t) + + httpID, err := st.SaveRepeaterTab(makeTab("HTTP", "example.com", 80, false)) + require.NoError(t, err) + wsID, err := st.SaveRepeaterTab(makeWSTab("WS", "wss://example.com/ws", true)) + require.NoError(t, err) + + require.NoError(t, st.DeleteRepeaterTab(wsID)) + + tabs, err := st.AllRepeaterTabs() + require.NoError(t, err) + require.Len(t, tabs, 1) + assert.Equal(t, httpID, tabs[0].ID) + assert.Equal(t, "http", tabs[0].TabType) +} diff --git a/internal/store/scope.go b/internal/store/scope.go new file mode 100644 index 0000000..8c7c4a0 --- /dev/null +++ b/internal/store/scope.go @@ -0,0 +1,116 @@ +package store + +import ( + "fmt" + "net" + "strings" +) + +// ScopeEntry represents a single scope target. +type ScopeEntry struct { + ID int64 + Host string +} + +// AllScopeEntries returns all scope entries. +func (s *Store) AllScopeEntries() ([]ScopeEntry, error) { + rows, err := s.db.Query(`SELECT id, host FROM scope ORDER BY id ASC`) + if err != nil { + return nil, fmt.Errorf("store: query scope: %w", err) + } + defer rows.Close() + + var entries []ScopeEntry + for rows.Next() { + var e ScopeEntry + if err := rows.Scan(&e.ID, &e.Host); err != nil { + return nil, fmt.Errorf("store: scan scope: %w", err) + } + entries = append(entries, e) + } + return entries, rows.Err() +} + +// AddScopeEntry adds a new host to scope. +func (s *Store) AddScopeEntry(host string) error { + err := s.write(func() error { + _, err := s.db.Exec(`INSERT INTO scope (host) VALUES (?)`, host) + if err != nil { + return fmt.Errorf("store: add scope entry: %w", err) + } + return nil + }) + if err == nil { + s.invalidateScopeCache() + } + return err +} + +// DeleteScopeEntry removes a scope entry by ID. +func (s *Store) DeleteScopeEntry(id int64) error { + err := s.write(func() error { + _, err := s.db.Exec(`DELETE FROM scope WHERE id = ?`, id) + if err != nil { + return fmt.Errorf("store: delete scope entry: %w", err) + } + return nil + }) + if err == nil { + s.invalidateScopeCache() + } + return err +} + +// InScope returns true if the given host matches any scope entry. +// example.com matches example.com and all subdomains. +// Uses an in-memory cache populated on first call and invalidated on scope writes. +func (s *Store) InScope(host string) bool { + // Strip port if present; net.SplitHostPort handles IPv6 brackets correctly. + if h, _, err := net.SplitHostPort(host); err == nil { + host = h + } + + s.scopeCacheMu.RLock() + if s.scopeCacheSet { + entries := s.scopeCache + s.scopeCacheMu.RUnlock() + return matchesScope(entries, host) + } + s.scopeCacheMu.RUnlock() + + // Cache miss — populate under write lock. + s.scopeCacheMu.Lock() + // Re-check after acquiring write lock — another goroutine may have populated it. + if !s.scopeCacheSet { + entries, err := s.AllScopeEntries() + if err == nil { + s.scopeCache = entries + s.scopeCacheSet = true + } + } + entries := s.scopeCache + s.scopeCacheMu.Unlock() + + return matchesScope(entries, host) +} + +// invalidateScopeCache clears the scope cache so the next InScope call reloads from DB. +func (s *Store) invalidateScopeCache() { + s.scopeCacheMu.Lock() + s.scopeCache = nil + s.scopeCacheSet = false + s.scopeCacheMu.Unlock() +} + +// matchesScope returns true if host matches any entry in entries. +func matchesScope(entries []ScopeEntry, host string) bool { + if len(entries) == 0 { + return true // no scope defined, everything is in scope + } + for _, e := range entries { + if host == e.Host || strings.HasSuffix(host, "."+e.Host) { + return true + } + } + return false +} diff --git a/internal/store/scope_test.go b/internal/store/scope_test.go new file mode 100644 index 0000000..b1267c8 --- /dev/null +++ b/internal/store/scope_test.go @@ -0,0 +1,135 @@ +package store_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestAddScopeEntry_AndRetrieve(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 1) + assert.Equal(t, "example.com", entries[0].Host) + assert.Positive(t, entries[0].ID) +} + +func TestAddScopeEntry_Multiple(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + require.NoError(t, st.AddScopeEntry("other.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + assert.Len(t, entries, 2) +} + +func TestAddScopeEntry_Duplicate_Errors(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + err := st.AddScopeEntry("example.com") + assert.Error(t, err, "duplicate host must be rejected by UNIQUE constraint") +} + +func TestDeleteScopeEntry(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("example.com")) + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 1) + + require.NoError(t, st.DeleteScopeEntry(entries[0].ID)) + + entries, err = st.AllScopeEntries() + require.NoError(t, err) + assert.Empty(t, entries) +} + +func TestDeleteScopeEntry_NonExistent(t *testing.T) { + st := newTestStore(t) + assert.NoError(t, st.DeleteScopeEntry(9999)) +} + +func TestAllScopeEntries_OrderedByIDAsc(t *testing.T) { + st := newTestStore(t) + + require.NoError(t, st.AddScopeEntry("alpha.com")) + require.NoError(t, st.AddScopeEntry("beta.com")) + require.NoError(t, st.AddScopeEntry("gamma.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.Len(t, entries, 3) + assert.Equal(t, "alpha.com", entries[0].Host) + assert.Equal(t, "beta.com", entries[1].Host) + assert.Equal(t, "gamma.com", entries[2].Host) +} + +func TestInScope_EmptyScopeAllowsEverything(t *testing.T) { + st := newTestStore(t) + assert.True(t, st.InScope("anything.com")) + assert.True(t, st.InScope("other.net")) +} + +func TestInScope_ExactMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("example.com")) +} + +func TestInScope_SubdomainMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("sub.example.com")) + assert.True(t, st.InScope("deep.sub.example.com")) +} + +func TestInScope_NoMatch(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.False(t, st.InScope("other.com")) + assert.False(t, st.InScope("notexample.com")) + assert.False(t, st.InScope("example.com.evil.com")) +} + +func TestInScope_StripsPort(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + assert.True(t, st.InScope("example.com:8080")) + assert.True(t, st.InScope("example.com:443")) +} + +func TestInScope_AfterDeleteScopeBecomesOpen(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + + entries, err := st.AllScopeEntries() + require.NoError(t, err) + require.NoError(t, st.DeleteScopeEntry(entries[0].ID)) + + // scope is now empty — everything is in scope + assert.True(t, st.InScope("example.com")) + assert.True(t, st.InScope("completely-unrelated.com")) +} + +func TestInScope_MultipleEntries(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.AddScopeEntry("example.com")) + require.NoError(t, st.AddScopeEntry("target.io")) + + assert.True(t, st.InScope("example.com")) + assert.True(t, st.InScope("api.target.io")) + assert.False(t, st.InScope("other.com")) +} diff --git a/internal/store/search_test.go b/internal/store/search_test.go new file mode 100644 index 0000000..be21b3e --- /dev/null +++ b/internal/store/search_test.go @@ -0,0 +1,105 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func logTxWithBodies(t *testing.T, st *store.Store, method, host, url string, status int, reqBody, respBody string) { + t.Helper() + tx := store.Transaction{ + Timestamp: time.Now().UTC(), + Host: host, + Method: method, + URL: url, + ReqBody: []byte(reqBody), + RespBody: []byte(respBody), + StatusCode: status, + InScope: true, + } + require.NoError(t, st.Log(tx)) +} + +func TestSearchTransactions_FindsInRespBody(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "secret_token_here") + logTxWithBodies(t, st, "GET", "b.com", "https://b.com/", 200, "", "nothing special") + + results, err := st.SearchTransactions("secret_token", false, true, false) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "a.com", results[0].Host) +} + +func TestSearchTransactions_FindsInReqBody(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "POST", "a.com", "https://a.com/login", 200, "password=hunter2", "ok") + logTxWithBodies(t, st, "GET", "b.com", "https://b.com/", 200, "", "ok") + + results, err := st.SearchTransactions("hunter2", true, false, false) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "a.com", results[0].Host) +} + +func TestSearchTransactions_EmptyTerm_ReturnsNil(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "data") + + results, err := st.SearchTransactions("", true, true, false) + require.NoError(t, err) + assert.Nil(t, results) +} + +func TestSearchTransactions_NoFieldsSelected_ReturnsNil(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "data") + + results, err := st.SearchTransactions("data", false, false, false) + require.NoError(t, err) + assert.Nil(t, results) +} + +func TestSearchTransactions_NoMatch_ReturnsEmpty(t *testing.T) { + st := newTestStore(t) + logTxWithBodies(t, st, "GET", "a.com", "https://a.com/", 200, "", "nothing here") + + results, err := st.SearchTransactions("xyzzy_not_found", false, true, false) + require.NoError(t, err) + assert.Empty(t, results) +} + +func TestTransactionsPage_MethodFilter(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/", 200))) + require.NoError(t, st.Log(makeTx("POST", "a.com", "https://a.com/login", 200))) + require.NoError(t, st.Log(makeTx("PUT", "a.com", "https://a.com/item", 200))) + + results, err := st.TransactionsPage(0, store.TransactionFilter{ + ShowOutScope: true, + Methods: []string{"POST"}, + }) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, "POST", results[0].Method) +} + +func TestTransactionsPage_StatusRangeFilter(t *testing.T) { + st := newTestStore(t) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/ok", 200))) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/redir", 302))) + require.NoError(t, st.Log(makeTx("GET", "a.com", "https://a.com/err", 404))) + + results, err := st.TransactionsPage(0, store.TransactionFilter{ + ShowOutScope: true, + StatusMin: 400, + StatusMax: 499, + }) + require.NoError(t, err) + require.Len(t, results, 1) + assert.Equal(t, 404, results[0].StatusCode) +} diff --git a/internal/store/store_test.go b/internal/store/store_test.go new file mode 100644 index 0000000..6b9d012 --- /dev/null +++ b/internal/store/store_test.go @@ -0,0 +1,22 @@ +package store_test + +import ( + "os" + "testing" + + "github.com/shiv/internal/store" + "github.com/stretchr/testify/require" +) + +func newTestStore(t *testing.T) *store.Store { + t.Helper() + f, err := os.CreateTemp("", "shiv-store-test-*.shiv") + require.NoError(t, err) + f.Close() + t.Cleanup(func() { os.Remove(f.Name()) }) + + st, err := store.Open(f.Name()) + require.NoError(t, err) + t.Cleanup(func() { st.Close() }) + return st +} diff --git a/internal/store/websocket.go b/internal/store/websocket.go new file mode 100644 index 0000000..562c487 --- /dev/null +++ b/internal/store/websocket.go @@ -0,0 +1,198 @@ +package store + +import ( + "fmt" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/logger" +) + +const wsFrameSizeLimit = 1 << 20 // 1 MB per frame + +// WebSocketFrame is a single frame captured during a WebSocket session. +type WebSocketFrame struct { + ID uint64 + ConnectionID uint64 // foreign key → websocket_connections.id + Timestamp time.Time + Direction events.WebSocketDirection + Opcode events.WebSocketOpcode + Payload []byte +} + +// WebSocketConnection represents the initial HTTP upgrade that started a +// WebSocket session. It is tied to the history transaction of the upgrade +// request so the UI can link from history to the frame list. +type WebSocketConnection struct { + ID uint64 + HistoryID uint64 + Host string + URL string + TLS bool + InScope bool + Timestamp time.Time + FrameCount int // populated on read, not stored +} + +// LogWebSocketConnection inserts a new WebSocket connection record and returns +// its ID. Called once when the upgrade handshake completes. +func (s *Store) LogWebSocketConnection(conn WebSocketConnection) (uint64, error) { + var id uint64 + err := s.write(func() error { + res, err := s.db.Exec(` + INSERT INTO websocket_connections (history_id, host, url, tls, in_scope, timestamp) + VALUES (?, ?, ?, ?, ?, ?)`, + conn.HistoryID, + conn.Host, + conn.URL, + boolToInt(conn.TLS), + boolToInt(conn.InScope), + conn.Timestamp.UTC().Format(time.RFC3339), + ) + if err != nil { + return fmt.Errorf("store: log websocket connection: %w", err) + } + lastID, err := res.LastInsertId() + if err != nil { + return fmt.Errorf("store: websocket connection last insert id: %w", err) + } + id = uint64(lastID) + return nil + }) + if err != nil { + return 0, err + } + + // Notify UI listeners — non-blocking, same pattern as Updates. + conn.ID = id + select { + case s.WebSocketConnections <- conn: + default: + } + + return id, nil +} + +// LogWebSocketFrame appends a single frame to an existing connection. +func (s *Store) LogWebSocketFrame(frame WebSocketFrame) error { + err := s.write(func() error { + _, err := s.db.Exec(` + INSERT INTO websocket_frames (connection_id, timestamp, direction, opcode, payload) + VALUES (?, ?, ?, ?, ?)`, + frame.ConnectionID, + frame.Timestamp.UTC().Format(time.RFC3339), + int(frame.Direction), + int(frame.Opcode), + frame.Payload, + ) + if err != nil { + return fmt.Errorf("store: log websocket frame: %w", err) + } + return nil + }) + if err != nil { + return err + } + + // Notify UI listeners — non-blocking. + select { + case s.WebSocketFrames <- frame: + default: + } + + return nil +} + +// AllWebSocketConnections returns all WebSocket connections ordered by id +// descending, with frame counts. +func (s *Store) AllWebSocketConnections() ([]WebSocketConnection, error) { + rows, err := s.db.Query(` + SELECT c.id, c.history_id, c.host, c.url, c.tls, c.in_scope, c.timestamp, + COUNT(f.id) as frame_count + FROM websocket_connections c + LEFT JOIN websocket_frames f ON f.connection_id = c.id + GROUP BY c.id + ORDER BY c.id DESC`) + if err != nil { + return nil, fmt.Errorf("store: query websocket connections: %w", err) + } + defer rows.Close() + + var conns []WebSocketConnection + for rows.Next() { + var c WebSocketConnection + var ts string + var tlsInt, scopeInt int + if err := rows.Scan(&c.ID, &c.HistoryID, &c.Host, &c.URL, &tlsInt, &scopeInt, &ts, &c.FrameCount); err != nil { + return nil, fmt.Errorf("store: scan websocket connection: %w", err) + } + c.Timestamp, _ = time.Parse(time.RFC3339, ts) + c.TLS = tlsInt == 1 + c.InScope = scopeInt == 1 + conns = append(conns, c) + } + return conns, rows.Err() +} + +// FramesForConnection returns all frames for a given connection ordered by id +// ascending (chronological). +func (s *Store) FramesForConnection(connectionID uint64) ([]WebSocketFrame, error) { + rows, err := s.db.Query(` + SELECT id, connection_id, timestamp, direction, opcode, payload + FROM websocket_frames + WHERE connection_id = ? + ORDER BY id ASC`, connectionID) + if err != nil { + return nil, fmt.Errorf("store: query websocket frames: %w", err) + } + defer rows.Close() + + var frames []WebSocketFrame + for rows.Next() { + var f WebSocketFrame + var ts string + if err := rows.Scan(&f.ID, &f.ConnectionID, &ts, &f.Direction, &f.Opcode, &f.Payload); err != nil { + return nil, fmt.Errorf("store: scan websocket frame: %w", err) + } + f.Timestamp, _ = time.Parse(time.RFC3339, ts) + frames = append(frames, f) + } + return frames, rows.Err() +} + +// ObserveWebSocketConnection implements events.WebSocketConnectionObserver. +func (s *Store) ObserveWebSocketConnection(e events.WebSocketConnectionEvent) uint64 { + id, err := s.LogWebSocketConnection(WebSocketConnection{ + Host: e.Host, + URL: e.URL, + TLS: e.TLS, + InScope: s.InScope(e.Host), + Timestamp: e.Timestamp, + }) + if err != nil { + logger.Error("store: observe ws connection: %v", err) + return 0 + } + return id +} + +// ObserveWebSocketFrame implements events.WebSocketFrameObserver. +func (s *Store) ObserveWebSocketFrame(e events.WebSocketFrameEvent) events.WebSocketFrameResult { + if e.ConnectionID == 0 { + return events.WebSocketFrameResult{Payload: e.Payload} + } + payload := e.Payload + if len(payload) > wsFrameSizeLimit { + payload = payload[:wsFrameSizeLimit] + } + if err := s.LogWebSocketFrame(WebSocketFrame{ + ConnectionID: e.ConnectionID, + Timestamp: e.Timestamp, + Direction: e.Direction, + Opcode: e.Opcode, + Payload: payload, + }); err != nil { + logger.Error("store: observe ws frame: %v", err) + } + return events.WebSocketFrameResult{Payload: e.Payload} +} diff --git a/internal/store/websocket_test.go b/internal/store/websocket_test.go new file mode 100644 index 0000000..6340080 --- /dev/null +++ b/internal/store/websocket_test.go @@ -0,0 +1,200 @@ +package store_test + +import ( + "testing" + "time" + + "github.com/shiv/internal/events" + "github.com/shiv/internal/store" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestLogWebSocketConnection_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + conn := store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/chat", + TLS: true, + InScope: true, + Timestamp: time.Now().UTC().Truncate(time.Second), + } + id, err := st.LogWebSocketConnection(conn) + require.NoError(t, err) + assert.NotZero(t, id) + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 1) + assert.Equal(t, id, conns[0].ID) + assert.Equal(t, "example.com:443", conns[0].Host) + assert.Equal(t, "wss://example.com:443/chat", conns[0].URL) + assert.True(t, conns[0].TLS) + assert.True(t, conns[0].InScope) +} + +func TestLogWebSocketFrame_PersistsAllFields(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + TLS: true, + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + frame := store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC().Truncate(time.Second), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: []byte("hello world"), + } + require.NoError(t, st.LogWebSocketFrame(frame)) + + frames, err := st.FramesForConnection(connID) + require.NoError(t, err) + require.Len(t, frames, 1) + assert.Equal(t, connID, frames[0].ConnectionID) + assert.Equal(t, events.WebSocketClient, frames[0].Direction) + assert.Equal(t, events.WebSocketText, frames[0].Opcode) + assert.Equal(t, []byte("hello world"), frames[0].Payload) +} + +func TestAllWebSocketConnections_OrderedByIDDesc(t *testing.T) { + st := newTestStore(t) + + for _, host := range []string{"first.com:443", "second.com:443", "third.com:443"} { + _, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: host, + URL: "wss://" + host + "/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + } + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 3) + assert.Equal(t, "third.com:443", conns[0].Host) + assert.Equal(t, "second.com:443", conns[1].Host) + assert.Equal(t, "first.com:443", conns[2].Host) +} + +func TestAllWebSocketConnections_FrameCountIsCorrect(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + for i := 0; i < 5; i++ { + require.NoError(t, st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: events.WebSocketClient, + Opcode: events.WebSocketText, + Payload: []byte("msg"), + })) + } + + conns, err := st.AllWebSocketConnections() + require.NoError(t, err) + require.Len(t, conns, 1) + assert.Equal(t, 5, conns[0].FrameCount) +} + +func TestFramesForConnection_OrderedByIDAsc(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + + directions := []events.WebSocketDirection{ + events.WebSocketClient, + events.WebSocketServer, + events.WebSocketClient, + } + for _, dir := range directions { + require.NoError(t, st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: dir, + Opcode: events.WebSocketText, + Payload: []byte("x"), + })) + } + + frames, err := st.FramesForConnection(connID) + require.NoError(t, err) + require.Len(t, frames, 3) + assert.Equal(t, events.WebSocketClient, frames[0].Direction) + assert.Equal(t, events.WebSocketServer, frames[1].Direction) + assert.Equal(t, events.WebSocketClient, frames[2].Direction) +} + +func TestFramesForConnection_EmptyForUnknownConnection(t *testing.T) { + st := newTestStore(t) + frames, err := st.FramesForConnection(99999) + require.NoError(t, err) + assert.Empty(t, frames) +} + +func TestLogWebSocketConnection_EmitsToChannel(t *testing.T) { + st := newTestStore(t) + + conn := store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + } + _, err := st.LogWebSocketConnection(conn) + require.NoError(t, err) + + select { + case c := <-st.WebSocketConnections: + assert.Equal(t, "example.com:443", c.Host) + case <-time.After(time.Second): + t.Fatal("timed out waiting for WebSocketConnections channel emit") + } +} + +func TestLogWebSocketFrame_EmitsToChannel(t *testing.T) { + st := newTestStore(t) + + connID, err := st.LogWebSocketConnection(store.WebSocketConnection{ + Host: "example.com:443", + URL: "wss://example.com:443/", + Timestamp: time.Now().UTC(), + }) + require.NoError(t, err) + // Drain the connection channel. + <-st.WebSocketConnections + + err = st.LogWebSocketFrame(store.WebSocketFrame{ + ConnectionID: connID, + Timestamp: time.Now().UTC(), + Direction: events.WebSocketServer, + Opcode: events.WebSocketText, + Payload: []byte("streamed"), + }) + require.NoError(t, err) + + select { + case f := <-st.WebSocketFrames: + assert.Equal(t, connID, f.ConnectionID) + assert.Equal(t, events.WebSocketServer, f.Direction) + assert.Equal(t, []byte("streamed"), f.Payload) + case <-time.After(time.Second): + t.Fatal("timed out waiting for WebSocketFrames channel emit") + } +} |