agcodex_execpolicy/
policy_parser.rs

1#![allow(clippy::needless_lifetimes)]
2
3use crate::Opt;
4use crate::Policy;
5use crate::ProgramSpec;
6use crate::arg_matcher::ArgMatcher;
7use crate::opt::OptMeta;
8use log::info;
9use multimap::MultiMap;
10use regex_lite::Regex;
11use starlark::any::ProvidesStaticType;
12use starlark::environment::GlobalsBuilder;
13use starlark::environment::LibraryExtension;
14use starlark::environment::Module;
15use starlark::eval::Evaluator;
16use starlark::syntax::AstModule;
17use starlark::syntax::Dialect;
18use starlark::values::Heap;
19use starlark::values::list::UnpackList;
20use starlark::values::none::NoneType;
21use std::cell::RefCell;
22use std::collections::HashMap;
23
24pub struct PolicyParser {
25    policy_source: String,
26    unparsed_policy: String,
27}
28
29impl PolicyParser {
30    pub fn new(policy_source: &str, unparsed_policy: &str) -> Self {
31        Self {
32            policy_source: policy_source.to_string(),
33            unparsed_policy: unparsed_policy.to_string(),
34        }
35    }
36
37    pub fn parse(&self) -> starlark::Result<Policy> {
38        let mut dialect = Dialect::Extended.clone();
39        dialect.enable_f_strings = true;
40        let ast = AstModule::parse(&self.policy_source, self.unparsed_policy.clone(), &dialect)?;
41        let globals = GlobalsBuilder::extended_by(&[LibraryExtension::Typing])
42            .with(policy_builtins)
43            .build();
44        let module = Module::new();
45
46        let heap = Heap::new();
47
48        module.set("ARG_OPAQUE_VALUE", heap.alloc(ArgMatcher::OpaqueNonFile));
49        module.set("ARG_RFILE", heap.alloc(ArgMatcher::ReadableFile));
50        module.set("ARG_WFILE", heap.alloc(ArgMatcher::WriteableFile));
51        module.set("ARG_RFILES", heap.alloc(ArgMatcher::ReadableFiles));
52        module.set(
53            "ARG_RFILES_OR_CWD",
54            heap.alloc(ArgMatcher::ReadableFilesOrCwd),
55        );
56        module.set("ARG_POS_INT", heap.alloc(ArgMatcher::PositiveInteger));
57        module.set("ARG_SED_COMMAND", heap.alloc(ArgMatcher::SedCommand));
58        module.set(
59            "ARG_UNVERIFIED_VARARGS",
60            heap.alloc(ArgMatcher::UnverifiedVarargs),
61        );
62
63        let policy_builder = PolicyBuilder::new();
64        {
65            let mut eval = Evaluator::new(&module);
66            eval.extra = Some(&policy_builder);
67            eval.eval_module(ast, &globals)?;
68        }
69        let policy = policy_builder.build();
70        policy.map_err(|e| starlark::Error::new_kind(starlark::ErrorKind::Other(e.into())))
71    }
72}
73
74#[derive(Debug)]
75pub struct ForbiddenProgramRegex {
76    pub regex: regex_lite::Regex,
77    pub reason: String,
78}
79
80#[derive(Debug, ProvidesStaticType)]
81struct PolicyBuilder {
82    programs: RefCell<MultiMap<String, ProgramSpec>>,
83    forbidden_program_regexes: RefCell<Vec<ForbiddenProgramRegex>>,
84    forbidden_substrings: RefCell<Vec<String>>,
85}
86
87impl PolicyBuilder {
88    fn new() -> Self {
89        Self {
90            programs: RefCell::new(MultiMap::new()),
91            forbidden_program_regexes: RefCell::new(Vec::new()),
92            forbidden_substrings: RefCell::new(Vec::new()),
93        }
94    }
95
96    fn build(self) -> Result<Policy, regex_lite::Error> {
97        let programs = self.programs.into_inner();
98        let forbidden_program_regexes = self.forbidden_program_regexes.into_inner();
99        let forbidden_substrings = self.forbidden_substrings.into_inner();
100        Policy::new(programs, forbidden_program_regexes, forbidden_substrings)
101    }
102
103    fn add_program_spec(&self, program_spec: ProgramSpec) {
104        info!("adding program spec: {program_spec:?}");
105        let name = program_spec.program.clone();
106        let mut programs = self.programs.borrow_mut();
107        programs.insert(name.clone(), program_spec);
108    }
109
110    fn add_forbidden_substrings(&self, substrings: &[String]) {
111        let mut forbidden_substrings = self.forbidden_substrings.borrow_mut();
112        forbidden_substrings.extend_from_slice(substrings);
113    }
114
115    fn add_forbidden_program_regex(&self, regex: Regex, reason: String) {
116        let mut forbidden_program_regexes = self.forbidden_program_regexes.borrow_mut();
117        forbidden_program_regexes.push(ForbiddenProgramRegex { regex, reason });
118    }
119}
120
121#[starlark_module]
122fn policy_builtins(builder: &mut GlobalsBuilder) {
123    fn define_program<'v>(
124        program: String,
125        system_path: Option<UnpackList<String>>,
126        option_bundling: Option<bool>,
127        combined_format: Option<bool>,
128        options: Option<UnpackList<Opt>>,
129        args: Option<UnpackList<ArgMatcher>>,
130        forbidden: Option<String>,
131        should_match: Option<UnpackList<UnpackList<String>>>,
132        should_not_match: Option<UnpackList<UnpackList<String>>>,
133        eval: &mut Evaluator,
134    ) -> anyhow::Result<NoneType> {
135        let option_bundling = option_bundling.unwrap_or(false);
136        let system_path = system_path.map_or_else(Vec::new, |v| v.items.to_vec());
137        let combined_format = combined_format.unwrap_or(false);
138        let options = options.map_or_else(Vec::new, |v| v.items.to_vec());
139        let args = args.map_or_else(Vec::new, |v| v.items.to_vec());
140
141        let mut allowed_options = HashMap::<String, Opt>::new();
142        for opt in options {
143            let name = opt.name().to_string();
144            if allowed_options
145                .insert(opt.name().to_string(), opt)
146                .is_some()
147            {
148                return Err(anyhow::format_err!("duplicate flag: {name}"));
149            }
150        }
151
152        let program_spec = ProgramSpec::new(
153            program,
154            system_path,
155            option_bundling,
156            combined_format,
157            allowed_options,
158            args,
159            forbidden,
160            should_match
161                .map_or_else(Vec::new, |v| v.items.to_vec())
162                .into_iter()
163                .map(|v| v.items.to_vec())
164                .collect(),
165            should_not_match
166                .map_or_else(Vec::new, |v| v.items.to_vec())
167                .into_iter()
168                .map(|v| v.items.to_vec())
169                .collect(),
170        );
171
172        #[expect(clippy::unwrap_used)]
173        let policy_builder = eval
174            .extra
175            .as_ref()
176            .unwrap()
177            .downcast_ref::<PolicyBuilder>()
178            .unwrap();
179        policy_builder.add_program_spec(program_spec);
180        Ok(NoneType)
181    }
182
183    fn forbid_substrings(
184        strings: UnpackList<String>,
185        eval: &mut Evaluator,
186    ) -> anyhow::Result<NoneType> {
187        #[expect(clippy::unwrap_used)]
188        let policy_builder = eval
189            .extra
190            .as_ref()
191            .unwrap()
192            .downcast_ref::<PolicyBuilder>()
193            .unwrap();
194        policy_builder.add_forbidden_substrings(&strings.items.to_vec());
195        Ok(NoneType)
196    }
197
198    fn forbid_program_regex(
199        regex: String,
200        reason: String,
201        eval: &mut Evaluator,
202    ) -> anyhow::Result<NoneType> {
203        #[expect(clippy::unwrap_used)]
204        let policy_builder = eval
205            .extra
206            .as_ref()
207            .unwrap()
208            .downcast_ref::<PolicyBuilder>()
209            .unwrap();
210        let compiled_regex = regex_lite::Regex::new(&regex)?;
211        policy_builder.add_forbidden_program_regex(compiled_regex, reason);
212        Ok(NoneType)
213    }
214
215    fn opt(name: String, r#type: ArgMatcher, required: Option<bool>) -> anyhow::Result<Opt> {
216        Ok(Opt::new(
217            name,
218            OptMeta::Value(r#type.arg_type()),
219            required.unwrap_or(false),
220        ))
221    }
222
223    fn flag(name: String) -> anyhow::Result<Opt> {
224        Ok(Opt::new(name, OptMeta::Flag, false))
225    }
226}
agcodex_execpolicy/policy_parser.rs

agcodex_execpolicy/
policy_parser.rs
