summaryrefslogtreecommitdiff
path: root/internal/store/project.go
blob: e379f320cf17bc63dd58fb4e65afd94fb4963a12 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
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 ''
);
`