ctf_brute/ops/pattern/
iter.rs

1use std::ops::{Range, RangeInclusive};
2
3use super::{Pattern, PatternIter};
4
5impl Iterator for PatternIter {
6    type Item = String;
7
8    fn next(&mut self) -> Option<Self::Item> {
9        match self {
10            Self::Range(range) => Some(range.next()?.to_string()),
11            Self::MRange(range) => Some(range.next()?.to_string()),
12            Self::Empty(b) => {
13                if *b {
14                    None
15                } else {
16                    *b = true;
17                    Some(String::new())
18                }
19            }
20            Self::Group(patterns, iterators, last) => {
21                assert_eq!(patterns.len(), iterators.len());
22                assert_eq!(iterators.len(), last.len());
23                let range = Range {
24                    start: 0,
25                    end: iterators.len(),
26                };
27                let mut next = true;
28                let mut result = String::new();
29                for i in range.into_iter().rev() {
30                    if last[i].is_none() || next {
31                        let mut next_val = iterators[i].next();
32                        if i == 0 && next_val.is_none() {
33                            return None; //End
34                        }
35                        if next_val.is_none() {
36                            iterators[i] = patterns[i].iter();
37                            next_val = iterators[i].next();
38                            last[i] = next_val;
39                        } else {
40                            next = false;
41                            last[i] = next_val;
42                        }
43                    }
44                    //Error one of iterators had 0 results (shouldnt be possible)
45                    if last[i].is_none() {
46                        return None;
47                    }
48                    let res = last[i].as_ref().unwrap();
49                    result.insert_str(0, res.as_str());
50                }
51                Some(result)
52            }
53            Self::Length(iteratrs) => loop {
54                if iteratrs.is_empty() {
55                    return None;
56                }
57                let next = iteratrs[0].next();
58                if next.is_none() {
59                    iteratrs.remove(0);
60                    continue;
61                }
62                return next;
63            },
64        }
65    }
66}
67
68impl Pattern {
69    pub fn iter(&self) -> PatternIter {
70        match &self {
71            Self::Range(range) => PatternIter::Range(range.iter()),
72            Self::MRange(range) => PatternIter::MRange(range.iter()),
73            Self::Group(patterns) => PatternIter::Group(
74                patterns.clone(),
75                patterns.iter().map(Pattern::iter).collect(),
76                patterns.iter().map(|_| None).collect(),
77            ),
78            Self::Length(pattern, range) => {
79                let patterns: Vec<PatternIter> = range
80                    .clone()
81                    .into_iter()
82                    .map(|i| {
83                        if i == 0 {
84                            return Self::Empty().into_iter();
85                        }
86                        let patterns: Vec<Self> = RangeInclusive::new(1, i)
87                            .map(|_| *pattern.clone())
88                            .collect();
89                        Self::Group(patterns).into_iter()
90                    })
91                    .collect();
92                PatternIter::Length(patterns)
93            }
94            Self::Empty() => PatternIter::Empty(false),
95        }
96    }
97}
98
99impl IntoIterator for Pattern {
100    type Item = String;
101
102    type IntoIter = PatternIter;
103
104    fn into_iter(self) -> Self::IntoIter {
105        self.iter()
106    }
107}
108
109#[cfg(test)]
110mod tests {
111    use crate::ops::BruteRange;
112
113    use super::*;
114    use nonempty::{nonempty, NonEmpty};
115
116    #[test]
117    fn test_group() {
118        let pattern = Pattern::Group(vec![
119            Pattern::Range(BruteRange::from_range('a'..='b')),
120            Pattern::Range(BruteRange::from_range('a'..='c')),
121        ]);
122        let result: Vec<String> = pattern.iter().collect();
123        assert_eq!(result, vec!["aa", "ab", "ac", "ba", "bb", "bc"]);
124    }
125
126    #[test]
127    fn test_lenght() {
128        let pattern = Pattern::Length(
129            Box::new(Pattern::Range(BruteRange::from_range('a'..='c'))),
130            1..=2,
131        );
132        let result: Vec<String> = pattern.iter().collect();
133        assert_eq!(
134            result,
135            vec!["a", "b", "c", "aa", "ab", "ac", "ba", "bb", "bc", "ca", "cb", "cc"]
136        );
137    }
138
139    #[test]
140    fn test_lenght2() {
141        let pattern = Pattern::Length(
142            Box::new(Pattern::Range(BruteRange::from_range('a'..='c'))),
143            0..=3,
144        );
145        let result: Vec<String> = pattern.iter().collect();
146        assert_eq!(
147            result,
148            vec![
149                "", "a", "b", "c", "aa", "ab", "ac", "ba", "bb", "bc", "ca", "cb", "cc", "aaa",
150                "aab", "aac", "aba", "abb", "abc", "aca", "acb", "acc", "baa", "bab", "bac", "bba",
151                "bbb", "bbc", "bca", "bcb", "bcc", "caa", "cab", "cac", "cba", "cbb", "cbc", "cca",
152                "ccb", "ccc"
153            ]
154        );
155    }
156}
ctf_brute/ops/pattern/iter.rs

ctf_brute/ops/pattern/
iter.rs
