wafrift_evolution/evolution/
population.rs1use crate::lineage::Lineage;
2use rand::Rng;
3use serde::{Deserialize, Serialize};
4
5#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct Chromosome {
8 pub genes: Vec<(String, String)>,
10 pub fitness: f64,
12 pub evaluations: u32,
14 #[serde(default = "default_lineage")]
16 pub lineage: Lineage,
17}
18
19fn default_lineage() -> Lineage {
20 Lineage::genesis(0)
21}
22
23impl Chromosome {
24 #[must_use]
26 pub fn new(genes: Vec<(String, String)>) -> Self {
27 Self {
28 genes,
29 fitness: 0.0,
30 evaluations: 0,
31 lineage: Lineage::genesis(0),
32 }
33 }
34
35 #[must_use]
37 pub fn with_lineage(genes: Vec<(String, String)>, lineage: Lineage) -> Self {
38 Self {
39 genes,
40 fitness: 0.0,
41 evaluations: 0,
42 lineage,
43 }
44 }
45
46 pub fn record_verdict(&mut self, verdict: &crate::types::OracleVerdict) {
48 self.evaluations += 1;
49 let value = verdict.to_fitness();
50 let alpha = 2.0 / (f64::from(self.evaluations) + 1.0);
51 self.fitness = alpha * value + (1.0 - alpha) * self.fitness;
52 }
53
54 pub fn record(&mut self, passed: bool) {
56 self.record_verdict(&crate::types::OracleVerdict::from_bool(passed));
57 }
58
59 #[must_use]
61 pub fn gene(&self, name: &str) -> Option<&str> {
62 self.genes
63 .iter()
64 .find(|(gene_name, _)| gene_name == name)
65 .map(|(_, value)| value.as_str())
66 }
67
68 #[must_use]
70 pub fn has_gene(&self, name: &str) -> bool {
71 self.genes.iter().any(|(gene_name, _)| gene_name == name)
72 }
73
74 #[must_use]
76 pub fn active_gene_count(&self) -> usize {
77 self.genes
78 .iter()
79 .filter(|(_, value)| value != "None")
80 .count()
81 }
82
83 #[must_use]
85 pub fn hash(&self) -> u64 {
86 use std::collections::hash_map::DefaultHasher;
87 use std::hash::{Hash, Hasher};
88 let mut hasher = DefaultHasher::new();
89 for (name, value) in &self.genes {
90 name.hash(&mut hasher);
91 value.hash(&mut hasher);
92 }
93 hasher.finish()
94 }
95}
96
97#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct GenePool {
100 pub pools: Vec<(String, Vec<String>)>,
102}
103
104impl GenePool {
105 #[must_use]
107 pub fn default_wafrift() -> Self {
108 Self {
109 pools: vec![
110 (
111 "encoding".into(),
112 vec![
113 "None".into(),
114 "CaseAlternation".into(),
115 "UrlEncode".into(),
116 "DoubleUrlEncode".into(),
117 "TripleUrlEncode".into(),
118 "UnicodeEncode".into(),
119 "HtmlEntityEncode".into(),
120 "OverlongUtf8".into(),
121 "WhitespaceInsertion".into(),
122 "SqlCommentInsertion".into(),
123 "NullByteInsertion".into(),
124 "ChunkedSplit".into(),
125 "ParameterPollution".into(),
126 ],
127 ),
128 (
129 "content_type".into(),
130 vec![
131 "None".into(),
132 "Multipart".into(),
133 "MultipartQuotedBoundary".into(),
134 "JsonNested".into(),
135 "JsonUnicodeKeys".into(),
136 "JsonWithComments".into(),
137 "XmlCdata".into(),
138 "XmlNamespace".into(),
139 "MixedContentType".into(),
140 ],
141 ),
142 (
143 "header_obfuscation".into(),
144 vec![
145 "None".into(),
146 "CaseMixing".into(),
147 "TabSeparator".into(),
148 "WhitespacePadding".into(),
149 "LineFolding".into(),
150 "UnderscoreSubstitution".into(),
151 ],
152 ),
153 (
154 "grammar_rule".into(),
155 vec![
156 "None".into(),
157 "tautology_swap".into(),
158 "comment_swap".into(),
159 "whitespace_swap".into(),
160 "equality_swap".into(),
161 "union_swap".into(),
162 "string_split".into(),
163 "mysql_conditional".into(),
164 "tag_event_swap".into(),
165 "exec_fn_swap".into(),
166 "uri_scheme".into(),
167 "separator_swap".into(),
168 "command_obfuscate".into(),
169 "ifs_swap".into(),
170 "path_obfuscate".into(),
171 "variable_indirection".into(),
172 ],
173 ),
174 ],
175 }
176 }
177
178 #[must_use]
180 pub fn values_for(&self, gene_name: &str) -> Option<&[String]> {
181 self.pools
182 .iter()
183 .find(|(name, _)| name == gene_name)
184 .map(|(_, values)| values.as_slice())
185 }
186
187 #[must_use]
189 pub fn gene_names(&self) -> Vec<&str> {
190 self.pools.iter().map(|(name, _)| name.as_str()).collect()
191 }
192
193 #[must_use]
195 pub fn random_value(&self, gene_name: &str, rng: &mut impl Rng) -> Option<String> {
196 let values = self.values_for(gene_name)?;
197 if values.is_empty() {
198 return None;
199 }
200 Some(values[rng.gen_range(0..values.len())].clone())
201 }
202
203 #[must_use]
205 pub fn all_values(&self) -> Vec<String> {
206 let mut values = Vec::new();
207 for (_, pool_values) in &self.pools {
208 for v in pool_values {
209 if !values.contains(v) {
210 values.push(v.clone());
211 }
212 }
213 }
214 values
215 }
216}
217
218#[must_use]
220pub fn random_chromosome(gene_pool: &GenePool, rng: &mut impl Rng) -> Chromosome {
221 let genes = gene_pool
222 .gene_names()
223 .into_iter()
224 .map(|name| {
225 let value = gene_pool
226 .random_value(name, rng)
227 .unwrap_or_else(|| String::from("None"));
228 (name.to_string(), value)
229 })
230 .collect();
231 Chromosome::new(genes)
232}
233
234#[must_use]
236pub fn baseline_chromosome(gene_pool: &GenePool) -> Chromosome {
237 let genes = gene_pool
238 .gene_names()
239 .into_iter()
240 .map(|name| (name.to_string(), String::from("None")))
241 .collect();
242 Chromosome::new(genes)
243}