mirror of
https://github.com/SpecterOps/Nemesis
synced 2026-06-08 12:36:42 +00:00
Lee Christensen
397 lines
13 KiB
PL/PgSQL
397 lines
13 KiB
PL/PgSQL
CREATE EXTENSION IF NOT EXISTS pg_trgm;
|
|
|
|
-----------------------
|
|
-- FILES
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS files (
|
|
object_id UUID PRIMARY KEY,
|
|
agent_id VARCHAR(255),
|
|
project VARCHAR(255),
|
|
timestamp TIMESTAMP WITH TIME ZONE,
|
|
expiration TIMESTAMP WITH TIME ZONE,
|
|
path TEXT,
|
|
originating_object_id UUID,
|
|
nesting_level INTEGER,
|
|
file_creation_time TIMESTAMP WITH TIME ZONE,
|
|
file_access_time TIMESTAMP WITH TIME ZONE,
|
|
file_modification_time TIMESTAMP WITH TIME ZONE,
|
|
security_info JSONB,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
|
|
-----------------------
|
|
-- ENRICHED FILES
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS files_enriched (
|
|
object_id UUID PRIMARY KEY,
|
|
agent_id VARCHAR(255),
|
|
project VARCHAR(255),
|
|
timestamp TIMESTAMP WITH TIME ZONE,
|
|
expiration TIMESTAMP WITH TIME ZONE,
|
|
path TEXT,
|
|
file_name VARCHAR(255),
|
|
extension VARCHAR(50),
|
|
size BIGINT,
|
|
magic_type TEXT,
|
|
mime_type TEXT,
|
|
is_plaintext BOOLEAN,
|
|
is_container BOOLEAN,
|
|
originating_object_id UUID,
|
|
nesting_level INTEGER,
|
|
file_creation_time TIMESTAMP WITH TIME ZONE,
|
|
file_access_time TIMESTAMP WITH TIME ZONE,
|
|
file_modification_time TIMESTAMP WITH TIME ZONE,
|
|
security_info JSONB,
|
|
hashes JSONB,
|
|
file_tags TEXT[] DEFAULT '{}',
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS enrichments (
|
|
enrichment_id BIGSERIAL PRIMARY KEY,
|
|
object_id UUID NOT NULL,
|
|
module_name VARCHAR(255) NOT NULL,
|
|
result_data JSONB NOT NULL,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (object_id) REFERENCES files_enriched(object_id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS transforms (
|
|
transform_id BIGSERIAL PRIMARY KEY,
|
|
object_id UUID NOT NULL,
|
|
type VARCHAR(255) NOT NULL,
|
|
transform_object_id UUID NOT NULL,
|
|
metadata JSONB,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (object_id) REFERENCES files_enriched(object_id) ON DELETE CASCADE
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS file_tags (
|
|
tag_name VARCHAR(255) PRIMARY KEY
|
|
);
|
|
|
|
-- findings are a bit further down
|
|
|
|
CREATE TABLE IF NOT EXISTS files_view_history (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
object_id UUID NOT NULL,
|
|
username VARCHAR(255) NOT NULL,
|
|
automated BOOLEAN,
|
|
timestamp TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS files_feedback (
|
|
object_id UUID PRIMARY KEY,
|
|
username VARCHAR(255), -- last person who left the feedback
|
|
automated BOOLEAN,
|
|
missing_parser BOOLEAN,
|
|
missing_file_viewer BOOLEAN,
|
|
sensitive_info_not_detected BOOLEAN,
|
|
comments TEXT,
|
|
alert_sent BOOLEAN DEFAULT false,
|
|
timestamp TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
FOREIGN KEY (object_id) REFERENCES files_enriched(object_id) ON DELETE CASCADE
|
|
);
|
|
|
|
|
|
-----------------------
|
|
-- Extracted File Features
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS files_enriched_dataset (
|
|
object_id UUID PRIMARY KEY,
|
|
agent_id VARCHAR(255),
|
|
project VARCHAR(255),
|
|
timestamp TIMESTAMP WITH TIME ZONE,
|
|
expiration TIMESTAMP WITH TIME ZONE,
|
|
path TEXT,
|
|
file_creation_time TIMESTAMP WITH TIME ZONE,
|
|
file_access_time TIMESTAMP WITH TIME ZONE,
|
|
file_modification_time TIMESTAMP WITH TIME ZONE,
|
|
features_version VARCHAR(255),
|
|
individual_features JSONB, -- individual features
|
|
sibling_features JSONB, -- features based on files in the same folder (and same agent)
|
|
agent_population_features JSONB, -- population features segmented per agent_id
|
|
global_population_features JSONB, -- global population features
|
|
labels JSONB,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
-- Create a GiST index for path matching
|
|
CREATE INDEX idx_files_enriched_dataset_path_trgm ON files_enriched_dataset USING gist (path gist_trgm_ops);
|
|
|
|
-- Create indexes for finding siblings (composite index with agent_id and path)
|
|
CREATE INDEX idx_files_enriched_dataset_siblings ON files_enriched_dataset(agent_id, path);
|
|
|
|
-- helper for pulling out 'sibling' files
|
|
CREATE OR REPLACE FUNCTION get_parent_path(path text)
|
|
RETURNS text AS $$
|
|
BEGIN
|
|
-- Handle both Windows and Unix paths
|
|
RETURN regexp_replace(
|
|
path,
|
|
'[\\/][^\\/]*$',
|
|
''
|
|
);
|
|
END;
|
|
$$ LANGUAGE plpgsql IMMUTABLE;
|
|
-- Then you can find siblings with queries like:
|
|
-- SELECT * FROM files_enriched_dataset
|
|
-- WHERE agent_id = $1
|
|
-- AND get_parent_path(path) = get_parent_path($2);
|
|
|
|
|
|
-----------------------
|
|
-- FINDINGS
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS findings (
|
|
finding_id BIGSERIAL PRIMARY KEY,
|
|
finding_name VARCHAR(255) NOT NULL,
|
|
category VARCHAR(255) NOT NULL,
|
|
severity INTEGER NOT NULL,
|
|
object_id UUID NOT NULL,
|
|
origin_type VARCHAR(255) NOT NULL,
|
|
origin_name VARCHAR(255) NOT NULL,
|
|
raw_data JSONB NOT NULL,
|
|
data JSONB NOT NULL,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
triage_id BIGINT
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS findings_triage_history (
|
|
id BIGSERIAL PRIMARY KEY,
|
|
finding_id BIGINT NOT NULL,
|
|
username VARCHAR(255) NOT NULL,
|
|
automated BOOLEAN,
|
|
value VARCHAR(255) NOT NULL,
|
|
timestamp TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
|
|
-- Add foreign key constraints if they don't exist
|
|
DO $$ BEGIN
|
|
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'fk_files_view_history_object_id') THEN
|
|
ALTER TABLE files_view_history
|
|
ADD CONSTRAINT fk_files_view_history_object_id
|
|
FOREIGN KEY (object_id)
|
|
REFERENCES files_enriched(object_id)
|
|
ON DELETE CASCADE;
|
|
END IF;
|
|
|
|
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'fk_findings_triage_id') THEN
|
|
ALTER TABLE findings
|
|
ADD CONSTRAINT fk_findings_triage_id
|
|
FOREIGN KEY (triage_id)
|
|
REFERENCES findings_triage_history(id);
|
|
END IF;
|
|
|
|
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'fk_findings_object_id') THEN
|
|
ALTER TABLE findings
|
|
ADD CONSTRAINT fk_findings_object_id
|
|
FOREIGN KEY (object_id)
|
|
REFERENCES files_enriched(object_id)
|
|
ON DELETE CASCADE;
|
|
END IF;
|
|
|
|
IF NOT EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'fk_findings_triage_history_finding_id') THEN
|
|
ALTER TABLE findings_triage_history
|
|
ADD CONSTRAINT fk_findings_triage_history_finding_id
|
|
FOREIGN KEY (finding_id)
|
|
REFERENCES findings(finding_id)
|
|
ON DELETE CASCADE;
|
|
END IF;
|
|
END $$;
|
|
|
|
|
|
-- Create indexes if they don't exist
|
|
CREATE INDEX IF NOT EXISTS idx_files_enriched_agent_id ON files_enriched(agent_id);
|
|
CREATE INDEX IF NOT EXISTS idx_files_enriched_hashes ON files_enriched USING GIN (hashes);
|
|
CREATE INDEX IF NOT EXISTS idx_files_enriched_path_trgm ON files_enriched USING gist (path gist_trgm_ops);
|
|
CREATE INDEX IF NOT EXISTS idx_findings_data_gin ON findings USING GIN (data);
|
|
CREATE INDEX IF NOT EXISTS idx_files_view_history_username ON files_view_history(username);
|
|
CREATE INDEX IF NOT EXISTS idx_files_view_history_composite ON files_view_history(object_id, username, timestamp);
|
|
CREATE INDEX IF NOT EXISTS idx_files_enriched_dataset_agent_id ON files_enriched_dataset(agent_id);
|
|
|
|
|
|
|
|
-----------------------
|
|
-- PLAINTEXT CONTENT
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS plaintext_content (
|
|
object_id UUID REFERENCES files_enriched(object_id) ON DELETE CASCADE,
|
|
chunk_number INTEGER,
|
|
content TEXT,
|
|
-- Use simple configuration to preserve all terms without stemming
|
|
content_vector tsvector GENERATED ALWAYS AS (to_tsvector('simple', content)) STORED,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
PRIMARY KEY (object_id, chunk_number)
|
|
);
|
|
|
|
-- Create GIN index for full-text search
|
|
CREATE INDEX IF NOT EXISTS idx_plaintext_content_vector ON plaintext_content USING GIN (content_vector);
|
|
--- CREATE INDEX IF NOT EXISTS idx_plaintext_content_vector ON plaintext_content USING GIN (content_vector gin_trgm_ops); --- trigram-based searches, which can be useful for partial-word matching?
|
|
|
|
-----------------------
|
|
-- SEARCH VIEW
|
|
-----------------------
|
|
-- Drop any existing function/view/type
|
|
DROP FUNCTION IF EXISTS public.search_documents CASCADE;
|
|
DROP TYPE IF EXISTS public.search_result CASCADE;
|
|
DROP VIEW IF EXISTS public.document_search_results CASCADE;
|
|
|
|
-- Create the view
|
|
CREATE OR REPLACE VIEW public.document_search_results AS
|
|
SELECT
|
|
pc.object_id,
|
|
pc.chunk_number,
|
|
pc.content,
|
|
pc.content_vector,
|
|
ef.file_name,
|
|
ef.path,
|
|
ef.extension,
|
|
ef.project,
|
|
ef.agent_id,
|
|
ef."timestamp"::timestamp with time zone
|
|
FROM plaintext_content pc
|
|
JOIN files_enriched ef ON pc.object_id = ef.object_id;
|
|
|
|
-- Create a simple search function that uses the view
|
|
CREATE OR REPLACE FUNCTION public.search_documents(
|
|
search_query text,
|
|
path_pattern text DEFAULT NULL,
|
|
agent_pattern text DEFAULT NULL,
|
|
project_name text DEFAULT NULL,
|
|
start_date timestamp with time zone DEFAULT NULL,
|
|
end_date timestamp with time zone DEFAULT NULL,
|
|
max_results integer DEFAULT 100
|
|
) RETURNS SETOF document_search_results
|
|
STABLE
|
|
LANGUAGE sql
|
|
AS $$
|
|
WITH ranked_chunks AS (
|
|
SELECT *,
|
|
ROW_NUMBER() OVER (
|
|
PARTITION BY object_id
|
|
ORDER BY chunk_number ASC
|
|
) as rn
|
|
FROM document_search_results
|
|
WHERE
|
|
content_vector @@ plainto_tsquery('simple', search_query)
|
|
AND (path_pattern IS NULL OR path LIKE path_pattern)
|
|
AND (agent_pattern IS NULL OR agent_id LIKE agent_pattern)
|
|
AND (project_name IS NULL OR project = project_name)
|
|
AND (start_date IS NULL OR "timestamp" >= start_date)
|
|
AND (end_date IS NULL OR "timestamp" <= end_date)
|
|
)
|
|
SELECT
|
|
object_id,
|
|
chunk_number,
|
|
content,
|
|
content_vector,
|
|
file_name,
|
|
path,
|
|
extension,
|
|
project,
|
|
agent_id,
|
|
"timestamp"
|
|
FROM ranked_chunks
|
|
WHERE rn = 1
|
|
ORDER BY "timestamp" DESC
|
|
LIMIT max_results;
|
|
$$;
|
|
|
|
-- compression
|
|
ALTER TABLE plaintext_content
|
|
ALTER COLUMN content SET STORAGE EXTENDED;
|
|
|
|
-- Add work_mem adjustment for text search operations
|
|
SET work_mem = '256MB';
|
|
|
|
|
|
|
|
-----------------------
|
|
-- Yara
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS yara_rules (
|
|
name VARCHAR(255) NOT NULL PRIMARY KEY,
|
|
content TEXT NOT NULL,
|
|
source VARCHAR(255),
|
|
enabled BOOLEAN NOT NULL DEFAULT true,
|
|
alert_enabled BOOLEAN NOT NULL DEFAULT true,
|
|
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|
|
|
|
|
|
-----------------------
|
|
-- CREATE UPDATE FIELD TRIGGER
|
|
-----------------------
|
|
|
|
-- Updates the "updated_at" value for tables we update
|
|
CREATE OR REPLACE FUNCTION update_updated_at_column()
|
|
RETURNS TRIGGER AS $$
|
|
BEGIN
|
|
IF row(NEW.*) IS DISTINCT FROM row(OLD.*) THEN
|
|
NEW.updated_at = CURRENT_TIMESTAMP;
|
|
END IF;
|
|
RETURN NEW;
|
|
END;
|
|
$$ language 'plpgsql';
|
|
|
|
-- Create triggers for each table
|
|
CREATE OR REPLACE TRIGGER update_files_updated_at
|
|
BEFORE UPDATE ON files
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_files_enriched_updated_at
|
|
BEFORE UPDATE ON files_enriched
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_yara_rules_updated_at
|
|
BEFORE UPDATE ON yara_rules
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_enrichments_updated_at
|
|
BEFORE UPDATE ON enrichments
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_enrichments_updated_at
|
|
BEFORE UPDATE ON transforms
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_findings_updated_at
|
|
BEFORE UPDATE ON findings
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
CREATE OR REPLACE TRIGGER update_files_enriched_dataset_updated_at
|
|
BEFORE UPDATE ON files_enriched_dataset
|
|
FOR EACH ROW
|
|
EXECUTE FUNCTION update_updated_at_column();
|
|
|
|
|
|
|
|
-----------------------
|
|
-- Workflow tracking
|
|
-----------------------
|
|
CREATE TABLE IF NOT EXISTS workflows (
|
|
wf_id VARCHAR(255) NOT NULL PRIMARY KEY,
|
|
object_id UUID,
|
|
filename VARCHAR(255),
|
|
enrichments_success TEXT[] DEFAULT '{}',
|
|
enrichments_failure TEXT[] DEFAULT '{}',
|
|
status TEXT NOT NULL,
|
|
runtime_seconds REAL,
|
|
start_time TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
|
|
);
|