// SPDX-License-Identifier: GPL-3.0-or-later

// Package correlation adds higher-confidence stories without hiding or
// replacing their raw alerts. It is transport- and persistence-independent.
package correlation

import "sort"

const SIDMultiStageIntrusion = 100080

type Incident struct {
	FirstTimestamp float64  `json:"first_ts"`
	LastTimestamp  float64  `json:"last_ts"`
	Signatures     []string `json:"signatures"`
}

type Record struct {
	SID               int      `json:"sid"`
	Signature         string   `json:"signature"`
	Classtype         string   `json:"classtype"`
	Severity          string   `json:"severity"`
	Summary           string   `json:"summary"`
	WindowSeconds     int      `json:"window_seconds"`
	MatchedSignatures []string `json:"matched_signatures"`
}

type Rule struct {
	SID         int
	Signature   string
	Classtype   string
	Severity    string
	Summary     string
	RequiredAny []map[string]bool
	MaxWindow   int
}

var DefaultRules = []Rule{{
	SID: SIDMultiStageIntrusion, Signature: "correlation.multi_stage_intrusion",
	Classtype: "multi-stage-intrusion", Severity: "CRITICAL",
	Summary: "Web/database service spawned a shell and the same incident showed " +
		"suspicious egress or an unusual listener within the correlation window",
	RequiredAny: []map[string]bool{
		{"exec_rule.web-rce": true},
		{"host_rule.suspicious-egress": true, "host_rule.suspicious-listener": true},
	},
	MaxWindow: 600,
}}

func Correlate(incident Incident, rules []Rule) []Record {
	if rules == nil {
		rules = DefaultRules
	}
	signatures := make(map[string]bool, len(incident.Signatures))
	for _, signature := range incident.Signatures {
		signatures[signature] = true
	}
	duration := incident.LastTimestamp - incident.FirstTimestamp
	result := make([]Record, 0)
	for _, rule := range rules {
		if duration > float64(rule.MaxWindow) {
			continue
		}
		matched := true
		for _, choices := range rule.RequiredAny {
			found := false
			for choice := range choices {
				if signatures[choice] {
					found = true
					break
				}
			}
			if !found {
				matched = false
				break
			}
		}
		if !matched {
			continue
		}
		matchedSignatures := make([]string, 0, len(signatures))
		for signature := range signatures {
			matchedSignatures = append(matchedSignatures, signature)
		}
		sort.Strings(matchedSignatures)
		result = append(result, Record{
			SID: rule.SID, Signature: rule.Signature, Classtype: rule.Classtype,
			Severity: rule.Severity, Summary: rule.Summary,
			WindowSeconds: rule.MaxWindow, MatchedSignatures: matchedSignatures,
		})
	}
	return result
}