spreadsheet_mcp/tools/
vba.rs

1use crate::model::{
2    VbaModuleDescriptor, VbaModuleSourceResponse, VbaProjectSummaryResponse,
3    VbaReferenceDescriptor, WorkbookId,
4};
5use crate::state::AppState;
6use anyhow::{Result, anyhow, bail};
7use schemars::JsonSchema;
8use serde::Deserialize;
9use std::fs::File;
10use std::io::Read;
11use std::path::Path;
12use std::sync::Arc;
13use zip::result::ZipError;
14
15const MAX_VBA_PROJECT_BYTES: u64 = 20 * 1024 * 1024;
16const DEFAULT_MAX_MODULES: u32 = 200;
17const DEFAULT_INCLUDE_REFERENCES: bool = true;
18
19const DEFAULT_OFFSET_LINES: u32 = 0;
20const DEFAULT_LIMIT_LINES: u32 = 200;
21const MAX_LIMIT_LINES: u32 = 5_000;
22
23#[derive(Debug, Deserialize, JsonSchema)]
24pub struct VbaProjectSummaryParams {
25    #[serde(alias = "workbook_id")]
26    pub workbook_or_fork_id: WorkbookId,
27    #[serde(default)]
28    pub max_modules: Option<u32>,
29    #[serde(default)]
30    pub include_references: Option<bool>,
31}
32
33pub async fn vba_project_summary(
34    state: Arc<AppState>,
35    params: VbaProjectSummaryParams,
36) -> Result<VbaProjectSummaryResponse> {
37    let workbook = state.open_workbook(&params.workbook_or_fork_id).await?;
38    let raw = extract_vba_project_bin(&workbook.path)?;
39
40    if raw.is_none() {
41        return Ok(VbaProjectSummaryResponse {
42            workbook_id: workbook.id.clone(),
43            workbook_short_id: workbook.short_id.clone(),
44            has_vba: false,
45            code_page: None,
46            sys_kind: None,
47            modules: Vec::new(),
48            modules_truncated: false,
49            references: Vec::new(),
50            references_truncated: false,
51            notes: vec!["No xl/vbaProject.bin found in workbook".to_string()],
52        });
53    }
54
55    let project = ovba::open_project(raw.unwrap())?;
56
57    let max_modules = params.max_modules.unwrap_or(DEFAULT_MAX_MODULES).max(1);
58    let include_references = params
59        .include_references
60        .unwrap_or(DEFAULT_INCLUDE_REFERENCES);
61
62    let mut modules: Vec<VbaModuleDescriptor> = Vec::new();
63    for module in project.modules.iter().take(max_modules as usize) {
64        let module_type = match module.module_type {
65            ovba::ModuleType::Procedural => "procedural",
66            ovba::ModuleType::DocClsDesigner => "doc_cls_designer",
67        }
68        .to_string();
69
70        modules.push(VbaModuleDescriptor {
71            name: module.name.clone(),
72            stream_name: module.stream_name.clone(),
73            doc_string: module.doc_string.clone(),
74            text_offset: module.text_offset as u64,
75            help_context: module.help_context,
76            module_type,
77            read_only: module.read_only,
78            private: module.private,
79        });
80    }
81
82    let modules_truncated = project.modules.len() > max_modules as usize;
83
84    let mut references: Vec<VbaReferenceDescriptor> = Vec::new();
85    let mut references_truncated = false;
86    if include_references {
87        for reference in project.references.iter() {
88            let (kind, debug) = summarize_reference(reference);
89            references.push(VbaReferenceDescriptor { kind, debug });
90            if references.len() >= 200 {
91                references_truncated = project.references.len() > references.len();
92                break;
93            }
94        }
95    }
96
97    let sys_kind = Some(
98        match project.information.sys_kind {
99            ovba::SysKind::Win16 => "win16",
100            ovba::SysKind::Win32 => "win32",
101            ovba::SysKind::MacOs => "macos",
102            ovba::SysKind::Win64 => "win64",
103        }
104        .to_string(),
105    );
106
107    Ok(VbaProjectSummaryResponse {
108        workbook_id: workbook.id.clone(),
109        workbook_short_id: workbook.short_id.clone(),
110        has_vba: true,
111        code_page: Some(project.information.code_page),
112        sys_kind,
113        modules,
114        modules_truncated,
115        references,
116        references_truncated,
117        notes: Vec::new(),
118    })
119}
120
121#[derive(Debug, Deserialize, JsonSchema)]
122pub struct VbaModuleSourceParams {
123    #[serde(alias = "workbook_id")]
124    pub workbook_or_fork_id: WorkbookId,
125    pub module_name: String,
126    #[serde(default = "default_offset_lines")]
127    pub offset_lines: u32,
128    #[serde(default = "default_limit_lines")]
129    pub limit_lines: u32,
130}
131
132fn default_offset_lines() -> u32 {
133    DEFAULT_OFFSET_LINES
134}
135
136fn default_limit_lines() -> u32 {
137    DEFAULT_LIMIT_LINES
138}
139
140pub async fn vba_module_source(
141    state: Arc<AppState>,
142    params: VbaModuleSourceParams,
143) -> Result<VbaModuleSourceResponse> {
144    let workbook = state.open_workbook(&params.workbook_or_fork_id).await?;
145    let raw = extract_vba_project_bin(&workbook.path)?
146        .ok_or_else(|| anyhow!("No xl/vbaProject.bin found in workbook"))?;
147
148    let project = ovba::open_project(raw)?;
149    let source = project.module_source(&params.module_name)?;
150
151    let offset = params.offset_lines;
152    let limit = params.limit_lines.clamp(1, MAX_LIMIT_LINES);
153
154    let mut total_lines: u32 = 0;
155    let mut selected: Vec<&str> = Vec::new();
156
157    for (idx, line) in source.lines().enumerate() {
158        let idx = idx as u32;
159        total_lines = total_lines.saturating_add(1);
160        if idx < offset {
161            continue;
162        }
163        if selected.len() >= limit as usize {
164            continue;
165        }
166        selected.push(line);
167    }
168
169    if total_lines == 0 && !source.is_empty() {
170        total_lines = 1;
171    }
172
173    let truncated = total_lines.saturating_sub(offset) > limit;
174
175    let mut page = selected.join("\n");
176    if !page.is_empty() {
177        page.push('\n');
178    }
179
180    Ok(VbaModuleSourceResponse {
181        workbook_id: workbook.id.clone(),
182        workbook_short_id: workbook.short_id.clone(),
183        module_name: params.module_name,
184        offset_lines: offset,
185        limit_lines: limit,
186        total_lines,
187        truncated,
188        source: page,
189    })
190}
191
192fn extract_vba_project_bin(path: &Path) -> Result<Option<Vec<u8>>> {
193    let file = File::open(path)
194        .map_err(|e| anyhow!("failed to open workbook {}: {}", path.display(), e))?;
195
196    let mut archive = zip::ZipArchive::new(file)
197        .map_err(|e| anyhow!("failed to open workbook zip {}: {}", path.display(), e))?;
198
199    let mut entry = match archive.by_name("xl/vbaProject.bin") {
200        Ok(f) => f,
201        Err(ZipError::FileNotFound) => return Ok(None),
202        Err(e) => return Err(anyhow!("failed to locate xl/vbaProject.bin: {}", e)),
203    };
204
205    let declared_size = entry.size();
206    if declared_size > MAX_VBA_PROJECT_BYTES {
207        bail!(
208            "xl/vbaProject.bin too large ({} bytes; max {} bytes)",
209            declared_size,
210            MAX_VBA_PROJECT_BYTES
211        );
212    }
213
214    let mut buf: Vec<u8> = Vec::with_capacity(declared_size.min(1024 * 1024) as usize);
215    entry
216        .read_to_end(&mut buf)
217        .map_err(|e| anyhow!("failed to read xl/vbaProject.bin: {}", e))?;
218
219    if buf.len() as u64 > MAX_VBA_PROJECT_BYTES {
220        bail!(
221            "xl/vbaProject.bin too large after read ({} bytes; max {} bytes)",
222            buf.len(),
223            MAX_VBA_PROJECT_BYTES
224        );
225    }
226
227    Ok(Some(buf))
228}
229
230fn summarize_reference(reference: &ovba::Reference) -> (String, String) {
231    let kind = match reference {
232        ovba::Reference::Control(_) => "control",
233        ovba::Reference::Original(_) => "original",
234        ovba::Reference::Registered(_) => "registered",
235        ovba::Reference::Project(_) => "project",
236    }
237    .to_string();
238
239    let mut debug = format!("{:?}", reference);
240    const MAX_DEBUG_BYTES: usize = 4096;
241    if debug.len() > MAX_DEBUG_BYTES {
242        debug.truncate(MAX_DEBUG_BYTES);
243        debug.push_str("...[truncated]");
244    }
245
246    (kind, debug)
247}