feat: Add multiple extractor support

GeekMasher · GeekMasher · commit fc6a3f63bd5b · 2025-04-11T15:50:27.000+01:00
diff --git a/action.yml b/action.yml
@@ -10,11 +10,10 @@ inputs:
   token:
     description: GitHub Token
   extractor:
-    description: GitHub Repository where the extractor is located
+    description: GitHub Repositories where the extractors are located
     required: true
   language:
     description: Language(s) to use
-    required: true
   attestation:
     description: Attestation
     default: 'false'
diff --git a/src/action.rs b/src/action.rs
@@ -25,15 +25,16 @@ pub struct Action {
     #[input(description = "GitHub Token")]
     token: String,
 
-    /// GitHub Repository where the extractor is located
+    /// GitHub Repositories where the extractors are located
     #[input(
-        description = "GitHub Repository where the extractor is located",
-        required = true
+        description = "GitHub Repositories where the extractors are located",
+        required = true,
+        split = ","
     )]
-    extractor: String,
+    extractor: Vec<String>,
 
-    /// Language(d) to use
-    #[input(description = "Language(s) to use", split = ",", required = true)]
+    /// Language(d) to use, e.g. `iac`, `javascript`, `python`, etc.
+    #[input(description = "Language(s) to use", split = ",")]
     language: Vec<String>,
 
     /// Attestation
@@ -49,19 +50,22 @@ pub struct Action {
 }
 
 impl Action {
-    /// Gets the repository to use for the extractor. If the repository is not provided,
+    /// Gets the repositories to use for the extractors. If no repositories are provided,
     /// it will use the repository that the action is running in.
-    pub fn extractor_repository(&self) -> Result<Repository> {
-        let repo = if self.extractor.is_empty() {
-            log::debug!("No extractor repository provided, using the current repository");
-            self.get_repository()?
-        } else {
-            log::debug!("Using the provided extractor repository");
-            self.extractor.clone()
-        };
-        log::info!("Extractor Repository :: {}", repo);
-
-        Ok(Repository::parse(&repo)?)
+    pub fn extractor_repositories(&self) -> Result<Vec<Repository>> {
+        let mut repositories = Vec::new();
+        for extractor in &self.extractor {
+            let repo = if extractor.is_empty() {
+                log::debug!("No extractor repository provided, using the current repository");
+                self.get_repository()?
+            } else {
+                log::debug!("Using the provided extractor repository");
+                extractor.clone()
+            };
+            log::info!("Extractor Repository :: {}", repo);
+            repositories.push(Repository::parse(&repo)?);
+        }
+        Ok(repositories)
     }
 
     pub fn languages(&self) -> Vec<CodeQLLanguage> {
@@ -107,12 +111,63 @@ mod tests {
 
     fn action() -> Action {
         Action {
-            extractor: "owner/repo".to_string(),
+            extractor: vec!["owner/repo1".to_string(), "owner/repo2".to_string()],
+            language: vec!["iac".to_string()],
+            ..Default::default()
+        }
+    }
+
+    fn action_with_extractors(extractors: Vec<String>) -> Action {
+        Action {
+            extractor: extractors,
             language: vec!["iac".to_string()],
             ..Default::default()
         }
     }
 
+    #[test]
+    fn test_extractor_repositories() {
+        let action = action();
+        let repositories = action.extractor_repositories().unwrap();
+        assert_eq!(repositories.len(), 2);
+        assert_eq!(repositories[0].to_string(), "owner/repo1");
+        assert_eq!(repositories[1].to_string(), "owner/repo2");
+    }
+
+    #[test]
+    fn test_extractor_repositories_multiple() {
+        let action = action_with_extractors(vec![
+            "owner/repo1".to_string(),
+            "owner/repo2".to_string(),
+        ]);
+        let repositories = action.extractor_repositories().unwrap();
+        assert_eq!(repositories.len(), 2);
+        assert_eq!(repositories[0].to_string(), "owner/repo1");
+        assert_eq!(repositories[1].to_string(), "owner/repo2");
+    }
+
+    #[test]
+    fn test_extractor_repositories_single() {
+        let action = action_with_extractors(vec!["owner/repo1".to_string()]);
+        let repositories = action.extractor_repositories().unwrap();
+        assert_eq!(repositories.len(), 1);
+        assert_eq!(repositories[0].to_string(), "owner/repo1");
+    }
+
+    #[test]
+    fn test_extractor_repositories_empty() {
+        let action = action_with_extractors(vec![]);
+        let result = action.extractor_repositories();
+        assert!(result.is_err(), "Expected error for empty extractor list");
+    }
+
+    #[test]
+    fn test_extractor_repositories_invalid_format() {
+        let action = action_with_extractors(vec!["invalid_repo_format".to_string()]);
+        let result = action.extractor_repositories();
+        assert!(result.is_err(), "Expected error for invalid repository format");
+    }
+
     #[test]
     fn test_validate_languages() {
         let action = action();
diff --git a/src/main.rs b/src/main.rs
@@ -3,7 +3,8 @@ use std::path::PathBuf;
 use anyhow::{Context, Result};
 use ghactions::{ActionTrait, ToolCache, group, groupend};
 use ghastoolkit::codeql::database::queries::CodeQLQueries;
-use ghastoolkit::{CodeQL, CodeQLDatabase};
+use ghastoolkit::codeql::CodeQLLanguage;
+use ghastoolkit::{CodeQL, CodeQLDatabase, CodeQLExtractor};
 use log::{debug, info};
 
 mod action;
@@ -24,8 +25,8 @@ async fn main() -> Result<()> {
     debug!("ToolCache :: {:?}", toolcache);
 
     // Extractor
-    let extractor_repo = action.extractor_repository()?;
-    info!("Extractor Repository :: {}", extractor_repo);
+    let extractor_repos = action.extractor_repositories()?;
+    info!("Extractor Repositories :: {:?}", extractor_repos);
 
     let extractor_path = PathBuf::from("./extractors");
     if !extractor_path.exists() {
@@ -34,21 +35,31 @@ async fn main() -> Result<()> {
         info!("Created Extractor Directory :: {:?}", extractor_path);
     }
 
-    let extractor = extractors::fetch_extractor(
-        &client,
-        &extractor_repo,
-        action.attestation(),
-        &extractor_path,
-    )
-    .await
-    .context("Failed to fetch extractor")?;
-    log::info!("Extractor :: {:?}", extractor);
-
-    let codeql = CodeQL::init()
-        .search_path(extractor)
-        .build()
+    let mut codeql_builder = CodeQL::init();
+    let mut extractors = Vec::new();
+
+    // Download and extract the extractor repositories
+    for extractor_repo in extractor_repos {
+        let extractor = extractors::fetch_extractor(
+            &client,
+            &extractor_repo,
+            action.attestation(),
+            &extractor_path,
+        )
         .await
-        .context("Failed to create CodeQL instance")?;
+        .context("Failed to fetch extractor")?;
+        log::info!("Extractor :: {:?}", extractor);
+
+        codeql_builder = codeql_builder.search_path(extractor_path.clone());
+
+        extractors.push((extractor_repo.clone(), CodeQLExtractor::load_path(extractor.clone())?));
+    }
+
+    let codeql = codeql_builder
+            .build()
+            .await
+            .context("Failed to create CodeQL instance")?;
+
     log::info!("CodeQL :: {:?}", codeql);
 
     let languages = codeql.get_languages().await?;
@@ -72,13 +83,22 @@ async fn main() -> Result<()> {
 
     std::fs::create_dir_all(&sarif_output)?;
 
-    for language in action.languages() {
-        let group = format!("Running {} extractor", language.language());
+    for (extractor_repo, extractor) in extractors.iter() {
+        let language = CodeQLLanguage::from(extractor.name.clone());
+
+        if !action.languages().is_empty() {
+            if !action.languages().contains(&language) {
+                log::info!("Skipping language :: {}", language);
+                continue;
+            }
+        }
+
+        let group = format!("Running `{}` extractor", language.language());
         group!(group);
 
         log::info!("Running extractor for language :: {}", language);
 
-        let database_path = databases.join(format!("db-{}", language));
+        let database_path = databases.join(format!("db-{}", language.language()));
         let sarif_path = sarif_output.join(format!("{}-results.sarif", language.language()));
 
         let database = CodeQLDatabase::init()
@@ -92,9 +112,10 @@ async fn main() -> Result<()> {
         codeql.database(&database).overwrite().create().await?;
         log::info!("Created database :: {:?}", database);
 
+        // TODO: Assumes the queries are in the same org
         let queries = CodeQLQueries::from(format!(
             "{}/{}-queries",
-            extractor_repo.owner.clone(),
+            extractor_repo.owner,
             language.language()
         ));
         log::debug!("Queries :: {:?}", queries);
