scrings/
python.rs

1use crate::error::Result;
2use crate::parser::Parser;
3use crate::rule::Rule;
4use crate::tree::{Node, Tree};
5use crate::visitor::LanguageVisitor;
6use std::cmp::{max, min};
7use tree_sitter_python::language as python_language;
8
9fn build_python_tree(source: &str) -> crate::error::Result<Tree> {
10    let mut parser = tree_sitter::Parser::new();
11    parser.set_language(&python_language()).unwrap();
12
13    let tree_sitter = parser.parse(source, None).unwrap();
14    Ok(Tree::new(source.as_bytes(), tree_sitter))
15}
16
17#[derive(Default)]
18pub struct Python;
19
20impl Parser for Python {
21    fn parse(&mut self, src: &str) -> crate::error::Result<Option<(u64, String)>> {
22        let tree = build_python_tree(src)?;
23
24        let mut detection_rule = (
25            LanguageVisitor::new(|c| {
26                matches!(
27                    c,
28                    "if_statement"
29                        | "for_statement"
30                        | "while_statement"
31                        | "try_statement"
32                        | "with_statement"
33                        | "function_definition"
34                        | "class_definition"
35                        | "decorated_definition"
36                        | "match_statement"
37                        | "future_import_statement"
38                        | "import_from_statement"
39                        | "assert_statement"
40                        | "raise_statement"
41                        | "pass_statement"
42                        | "exec_statement"
43                        | "import_statement"
44                )
45            }),
46            IsPythonSubscript::new(),
47            IsPythonFunction::new(),
48        );
49
50        tree.apply(&mut detection_rule)?;
51
52        let start = match (
53            detection_rule.0.start,
54            detection_rule.1.start,
55            detection_rule.2.start,
56        ) {
57            (None, None, None) => None,
58            (None, Some(x), None) | (Some(x), None, None) | (None, None, Some(x)) => Some(x),
59            (Some(x), Some(y), None) | (Some(x), None, Some(y)) | (None, Some(x), Some(y)) => {
60                Some(min(x, y))
61            }
62            (Some(x), Some(y), Some(z)) => Some(min(min(x, y), z)),
63        };
64
65        let end = match (
66            detection_rule.0.end,
67            detection_rule.1.end,
68            detection_rule.2.end,
69        ) {
70            (None, None, None) => None,
71            (None, Some(x), None) | (Some(x), None, None) | (None, None, Some(x)) => Some(x),
72            (Some(x), Some(y), None) | (Some(x), None, Some(y)) | (None, Some(x), Some(y)) => {
73                Some(max(x, y))
74            }
75            (Some(x), Some(y), Some(z)) => Some(max(max(x, y), z)),
76        };
77
78        Ok(
79            if detection_rule.0.is_matched
80                || detection_rule.1.is_subscript
81                || detection_rule.2.is_function
82            {
83                Some((
84                    start.unwrap_or(0) as u64,
85                    String::from(&src[start.unwrap_or(0)..end.unwrap_or(src.len())]),
86                ))
87            } else {
88                None
89            },
90        )
91    }
92}
93
94pub struct IsPythonSubscript {
95    is_subscript: bool,
96    start: Option<usize>,
97    end: Option<usize>,
98    stack: Vec<bool>,
99}
100
101impl IsPythonSubscript {
102    pub fn new() -> Self {
103        Self {
104            is_subscript: false,
105            start: None,
106            end: None,
107            stack: vec![true],
108        }
109    }
110
111    pub fn verify(node: &Node) -> bool {
112        match node.kind() {
113            "subscript" => {
114                if node.child_count() == 4
115                    && node.child(1).unwrap().kind() == "["
116                    && node.child(3).unwrap().kind() == "]"
117                {
118                    if let Some(right) = node.child(2) {
119                        if right.kind() == "slice" {
120                            return true;
121                        }
122                    }
123                }
124            }
125            _ => (),
126        }
127        false
128    }
129}
130
131impl<'a> Rule<'a> for IsPythonSubscript {
132    // Match python slice
133    // Verify if the parent is subscript
134    // Assert all the children of the parent are valid python nodes
135    fn enter(&mut self, node: &Node<'a>) -> Result<bool> {
136        if IsPythonSubscript::verify(node) {
137            self.stack.push(true);
138
139            self.start = Some(min(
140                self.start.unwrap_or(node.start_abs()),
141                node.start_abs(),
142            ));
143            self.end = Some(max(self.end.unwrap_or(node.end_abs()), node.end_abs()));
144        }
145        Ok(true)
146    }
147
148    fn leave(&mut self, node: &Node<'a>) -> Result<()> {
149        if node.kind() == "ERROR" || node.text()? == "" {
150            for c in self.stack.iter_mut() {
151                *c = false;
152            }
153        }
154
155        if node.child_count() > 1 {
156            if IsPythonSubscript::verify(node) {
157                if self.stack.pop().unwrap_or(false) {
158                    self.start = Some(min(
159                        self.start.unwrap_or(node.start_abs()),
160                        node.start_abs(),
161                    ));
162                    self.end = Some(max(self.end.unwrap_or(node.end_abs()), node.end_abs()));
163                    self.is_subscript = true;
164                }
165            }
166        }
167
168        if self.is_subscript && self.stack.last() == Some(&true) {
169            self.start = Some(min(
170                self.start.unwrap_or(node.start_abs()),
171                node.start_abs(),
172            ));
173            self.end = Some(max(self.end.unwrap_or(node.end_abs()), node.end_abs()));
174        }
175
176        Ok(())
177    }
178}
179
180pub struct IsPythonFunction {
181    is_function: bool,
182    start: Option<usize>,
183    end: Option<usize>,
184}
185
186impl IsPythonFunction {
187    pub fn new() -> Self {
188        Self {
189            is_function: false,
190            start: None,
191            end: None,
192        }
193    }
194}
195
196impl<'a> Rule<'a> for IsPythonFunction {
197    // Match python function if in list
198    fn enter(&mut self, node: &Node<'a>) -> Result<bool> {
199        match node.kind() {
200            "call" => {
201                if let Some(function) = node.named_child("function") {
202                    if matches!(
203                        function.text().unwrap_or(""),
204                        "requests.get"
205                            | "requests.post"
206                            | "os.system"
207                            | "base64.b64decode"
208                            | "b64decode"
209                            | "subprocess.run"
210                            | "subprocess.Popen"
211                            | "subprocess.call"
212                    ) {
213                        self.is_function = true;
214                        self.start = Some(min(
215                            self.start.unwrap_or(node.start_abs()),
216                            node.start_abs(),
217                        ));
218                        self.end = Some(max(self.end.unwrap_or(node.end_abs()), node.end_abs()));
219                    }
220                }
221            }
222            _ => (),
223        }
224
225        Ok(true)
226    }
227
228    fn leave(&mut self, _node: &Node<'a>) -> Result<()> {
229        Ok(())
230    }
231}
scrings/python.rs

scrings/
python.rs
