use std::collections::{BTreeMap, BTreeSet};
use std::sync::LazyLock;
use crate::{hex_encode, CredentialHash, MatchLocation, Severity, VerifiedFinding};
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum CorrelationKind {
ValueReuse,
SplitComposite,
}
impl CorrelationKind {
#[must_use]
pub fn as_str(self) -> &'static str {
match self {
Self::ValueReuse => "value_reuse",
Self::SplitComposite => "split_composite",
}
}
}
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize,
)]
#[serde(rename_all = "snake_case")]
pub enum CorrelationRole {
SameValue,
RequiredPart,
OptionalPart,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, serde::Serialize, serde::Deserialize)]
pub struct CorrelatedLocation {
pub file_path: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub line: Option<usize>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CorrelatedMember {
pub detector_id: String,
pub detector_name: String,
pub service: String,
pub severity: Severity,
pub role: CorrelationRole,
pub credential_redacted: String,
pub credential_hash: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub confidence: Option<f64>,
pub locations: Vec<CorrelatedLocation>,
}
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct CorrelatedCredential {
pub id: String,
pub kind: CorrelationKind,
pub title: String,
pub service: String,
pub severity: Severity,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub confidence: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub strongest_member_confidence: Option<f64>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub scope: Option<String>,
pub file_count: usize,
pub impact: String,
pub members: Vec<CorrelatedMember>,
pub locations: Vec<CorrelatedLocation>,
}
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct CorrelationSettings {
reuse_min_files: usize,
reuse_confidence_bonus: f64,
max_confidence: f64,
reuse_impact: String,
}
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct CompositeSpec {
id: String,
service: String,
name: String,
severity: Severity,
required: Vec<String>,
#[serde(default)]
optional: Vec<String>,
confidence_bonus: f64,
impact: String,
}
#[derive(serde::Deserialize)]
#[serde(deny_unknown_fields)]
struct CorrelationPolicy {
settings: CorrelationSettings,
#[serde(default)]
composite: Vec<CompositeSpec>,
}
#[allow(clippy::panic)]
static POLICY: LazyLock<CorrelationPolicy> = LazyLock::new(|| {
match parse_policy(
include_str!("../data/credential-correlation.toml"),
"<embedded data/credential-correlation.toml>",
) {
Ok(policy) => policy,
Err(error) => panic!(
"keyhog: credential-correlation policy '<embedded \
data/credential-correlation.toml>' is invalid: {error}. \
Fix: correct crates/core/data/credential-correlation.toml and rebuild"
),
}
});
fn parse_policy(raw: &str, origin: &str) -> Result<CorrelationPolicy, String> {
let policy = toml::from_str::<CorrelationPolicy>(raw)
.map_err(|error| format!("failed to parse {origin}: {error}"))?;
validate_policy(&policy, origin)?;
Ok(policy)
}
fn validate_policy(policy: &CorrelationPolicy, origin: &str) -> Result<(), String> {
let settings = &policy.settings;
if settings.reuse_min_files < 2 {
return Err(format!(
"{origin} [settings] reuse_min_files must be at least 2, got {}",
settings.reuse_min_files
));
}
if !(settings.reuse_confidence_bonus > 0.0 && settings.reuse_confidence_bonus <= 1.0) {
return Err(format!(
"{origin} [settings] reuse_confidence_bonus must be in (0.0, 1.0], got {}",
settings.reuse_confidence_bonus
));
}
if !(settings.max_confidence > 0.0 && settings.max_confidence <= 1.0) {
return Err(format!(
"{origin} [settings] max_confidence must be in (0.0, 1.0], got {}",
settings.max_confidence
));
}
if settings.reuse_impact.trim().is_empty() {
return Err(format!(
"{origin} [settings] reuse_impact must not be empty"
));
}
let mut seen_ids = BTreeSet::new();
for composite in &policy.composite {
let id = composite.id.trim();
if id.is_empty() {
return Err(format!("{origin} [[composite]] has an empty id"));
}
if !seen_ids.insert(id) {
return Err(format!("{origin} [[composite]] duplicate id {id:?}"));
}
if composite.service.trim().is_empty() {
return Err(format!(
"{origin} [[composite]] {id:?} has an empty service"
));
}
if composite.name.trim().is_empty() {
return Err(format!("{origin} [[composite]] {id:?} has an empty name"));
}
if composite.impact.trim().is_empty() {
return Err(format!("{origin} [[composite]] {id:?} has an empty impact"));
}
if composite.required.len() < 2 {
return Err(format!(
"{origin} [[composite]] {id:?} needs at least 2 required parts, got {}",
composite.required.len()
));
}
if !(composite.confidence_bonus > 0.0 && composite.confidence_bonus <= 1.0) {
return Err(format!(
"{origin} [[composite]] {id:?} confidence_bonus must be in (0.0, 1.0], got {}",
composite.confidence_bonus
));
}
let mut seen_parts = BTreeSet::new();
for part in composite.required.iter().chain(composite.optional.iter()) {
let part = part.trim();
if part.is_empty() {
return Err(format!(
"{origin} [[composite]] {id:?} has an empty part id"
));
}
if !seen_parts.insert(part) {
return Err(format!(
"{origin} [[composite]] {id:?} lists part {part:?} more than once"
));
}
}
}
Ok(())
}
#[must_use]
pub fn correlation_composite_part_ids() -> Vec<&'static str> {
let mut ids: Vec<&'static str> = POLICY
.composite
.iter()
.flat_map(|composite| composite.required.iter().chain(composite.optional.iter()))
.map(String::as_str)
.collect();
ids.sort_unstable();
ids.dedup();
ids
}
pub fn validate_correlation_policy(raw: &str, origin: &str) -> Result<(), String> {
parse_policy(raw, origin).map(|_| ())
}
fn parent_dir(path: &str) -> &str {
match path.rfind(['/', '\\']) {
Some(0) => &path[..1],
Some(index) => &path[..index],
None => ".",
}
}
fn finding_locations(finding: &VerifiedFinding) -> impl Iterator<Item = &MatchLocation> {
std::iter::once(&finding.location).chain(finding.additional_locations.iter())
}
fn strongest_confidence<'a>(members: impl Iterator<Item = &'a VerifiedFinding>) -> Option<f64> {
members
.filter_map(|finding| finding.confidence)
.fold(None, |best: Option<f64>, value| {
Some(best.map_or(value, |current| current.max(value)))
})
}
fn lift(strongest: Option<f64>, bonus: f64) -> Option<f64> {
strongest.map(|value| {
(value + bonus)
.min(POLICY.settings.max_confidence)
.max(value)
})
}
fn member_of(
finding: &VerifiedFinding,
role: CorrelationRole,
scope: Option<&str>,
) -> CorrelatedMember {
let mut locations: Vec<CorrelatedLocation> = finding_locations(finding)
.filter_map(|location| {
let path = location.file_path.as_deref()?;
if scope.is_some_and(|dir| parent_dir(path) != dir) {
return None;
}
Some(CorrelatedLocation {
file_path: path.to_string(),
line: location.line,
})
})
.collect();
locations.sort();
locations.dedup();
CorrelatedMember {
detector_id: finding.detector_id.to_string(),
detector_name: finding.detector_name.to_string(),
service: finding.service.to_string(),
severity: finding.severity,
role,
credential_redacted: finding.credential_redacted.to_string(),
credential_hash: hex_encode(finding.credential_hash),
confidence: finding.confidence,
locations,
}
}
fn union_locations(members: &[CorrelatedMember]) -> Vec<CorrelatedLocation> {
let mut locations: Vec<CorrelatedLocation> = members
.iter()
.flat_map(|member| member.locations.iter().cloned())
.collect();
locations.sort();
locations.dedup();
locations
}
fn distinct_files(locations: &[CorrelatedLocation]) -> usize {
locations
.iter()
.map(|location| location.file_path.as_str())
.collect::<BTreeSet<_>>()
.len()
}
fn shared_service(members: &[CorrelatedMember]) -> String {
let mut services = members.iter().map(|member| member.service.as_str());
let Some(first) = services.next() else {
return "multiple".to_string();
};
if services.all(|service| service == first) {
first.to_string()
} else {
"multiple".to_string()
}
}
#[must_use]
pub fn correlate_findings(findings: &[VerifiedFinding]) -> Vec<CorrelatedCredential> {
let mut correlations = value_reuse_groups(findings);
correlations.extend(split_composite_groups(findings));
correlations.sort_by(|left, right| {
left.kind
.cmp(&right.kind)
.then_with(|| right.severity.cmp(&left.severity))
.then_with(|| left.service.cmp(&right.service))
.then_with(|| left.id.cmp(&right.id))
});
correlations
}
fn value_reuse_groups(findings: &[VerifiedFinding]) -> Vec<CorrelatedCredential> {
let mut by_digest: BTreeMap<CredentialHash, Vec<&VerifiedFinding>> = BTreeMap::new();
for finding in findings {
by_digest
.entry(finding.credential_hash)
.or_default()
.push(finding);
}
let mut groups = Vec::new();
for (digest, mut group) in by_digest {
group.sort_by(|left, right| {
left.detector_id
.cmp(&right.detector_id)
.then_with(|| left.location.file_path.cmp(&right.location.file_path))
.then_with(|| left.location.line.cmp(&right.location.line))
});
let members: Vec<CorrelatedMember> = group
.iter()
.map(|finding| member_of(finding, CorrelationRole::SameValue, None))
.collect();
let locations = union_locations(&members);
let file_count = distinct_files(&locations);
if file_count < POLICY.settings.reuse_min_files {
continue;
}
let detectors: BTreeSet<&str> = members
.iter()
.map(|member| member.detector_id.as_str())
.collect();
let title = if detectors.len() > 1 {
format!(
"One secret value matched by {} detectors across {file_count} files",
detectors.len()
)
} else {
format!(
"{} value reused across {file_count} files",
members
.first()
.map_or("Credential", |member| member.detector_name.as_str())
)
};
let strongest = strongest_confidence(group.iter().copied());
let severity = members
.iter()
.map(|member| member.severity)
.max()
.unwrap_or_default(); groups.push(CorrelatedCredential {
id: format!("reuse:{}", hex_encode(digest)),
kind: CorrelationKind::ValueReuse,
title,
service: shared_service(&members),
severity,
confidence: lift(strongest, POLICY.settings.reuse_confidence_bonus),
strongest_member_confidence: strongest,
scope: None,
file_count,
impact: POLICY.settings.reuse_impact.clone(),
members,
locations,
});
}
groups
}
type DirectoryIndex<'a> = BTreeMap<&'a str, BTreeMap<&'a str, BTreeSet<CredentialHash>>>;
fn split_composite_groups(findings: &[VerifiedFinding]) -> Vec<CorrelatedCredential> {
let mut index: DirectoryIndex<'_> = BTreeMap::new();
let mut by_part: BTreeMap<(&str, CredentialHash), &VerifiedFinding> = BTreeMap::new();
for finding in findings {
by_part.insert((&finding.detector_id, finding.credential_hash), finding);
for location in finding_locations(finding) {
let Some(path) = location.file_path.as_deref() else {
continue;
};
index
.entry(parent_dir(path))
.or_default()
.entry(&finding.detector_id)
.or_default()
.insert(finding.credential_hash);
}
}
let mut groups = Vec::new();
for (directory, detectors) in &index {
for composite in &POLICY.composite {
let Some(group) = composite_group(composite, directory, detectors, &by_part) else {
continue;
};
groups.push(group);
}
}
groups
}
fn composite_group(
composite: &CompositeSpec,
directory: &str,
detectors: &BTreeMap<&str, BTreeSet<CredentialHash>>,
by_part: &BTreeMap<(&str, CredentialHash), &VerifiedFinding>,
) -> Option<CorrelatedCredential> {
let mut members = Vec::with_capacity(composite.required.len() + composite.optional.len());
let mut sources = Vec::with_capacity(composite.required.len());
for part in &composite.required {
let digests = detectors.get(part.as_str())?;
let [digest] = digests.iter().copied().collect::<Vec<_>>()[..] else {
return None;
};
let finding = by_part.get(&(part.as_str(), digest))?;
let member = member_of(finding, CorrelationRole::RequiredPart, Some(directory));
sources.push(*finding);
members.push(member);
}
let mut shared: Option<BTreeSet<&str>> = None;
for member in &members {
let files: BTreeSet<&str> = member
.locations
.iter()
.map(|location| location.file_path.as_str())
.collect();
shared = Some(match shared {
None => files,
Some(current) => current.intersection(&files).copied().collect(),
});
}
if shared.is_none_or(|files| !files.is_empty()) {
return None;
}
for part in &composite.optional {
let Some(digests) = detectors.get(part.as_str()) else {
continue;
};
let [digest] = digests.iter().copied().collect::<Vec<_>>()[..] else {
continue;
};
let Some(finding) = by_part.get(&(part.as_str(), digest)) else {
continue;
};
sources.push(*finding);
members.push(member_of(
finding,
CorrelationRole::OptionalPart,
Some(directory),
));
}
members.sort_by(|left, right| {
left.detector_id
.cmp(&right.detector_id)
.then_with(|| left.credential_hash.cmp(&right.credential_hash))
});
let locations = union_locations(&members);
let file_count = distinct_files(&locations);
let strongest = strongest_confidence(sources.into_iter());
let severity = members
.iter()
.map(|member| member.severity)
.max()
.unwrap_or_default() .max(composite.severity);
Some(CorrelatedCredential {
id: format!("composite:{}@{directory}", composite.id),
kind: CorrelationKind::SplitComposite,
title: format!("{} split across {file_count} files", composite.name),
service: composite.service.clone(),
severity,
confidence: lift(strongest, composite.confidence_bonus),
strongest_member_confidence: strongest,
scope: Some(directory.to_string()),
file_count,
impact: composite.impact.clone(),
members,
locations,
})
}