Add feature to redirect output to database (#25)

* Update .gitignore and add log.sh
Add db_service module to service/mod.rs
Add async to main function
Update tokio features in Cargo.toml
Add postgres and sea-orm dependencies to Cargo.toml
Add to_db field to ConfigDto struct in scan.rs
Update sensleaks function to aODsync
Update detect function to async
Add insert_leaks function to db_service.rs

Signed-off-by: yaokunzhang <[email protected]>

* Refactor insert_leaks function signature and simplify get_db_config function

Signed-off-by: yaokunzhang <[email protected]>

* Refactor get_db_config function

Signed-off-by: yaokunzhang <[email protected]>

---------

Signed-off-by: yaokunzhang <[email protected]>

Author: yaokunzhang

.gitignore

@@ -14,4 +14,6 @@ Cargo.lock
 *.pdb
 
 # Compiled binary file
-.gitignore
+.gitignore
+
+log.sh

Cargo.toml

@@ -28,8 +28,10 @@ csv = "1.1"
 log = "0.4"
 env_logger = "0.11.0"
 axum = { version = "0.7.4", features = ["macros"] }
-tokio = { version = "1.36.0", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.36.0", features = ["full"] }
 tower-http = { version = "0.5.0", features = ["cors"] }
 utoipa = { version = "4.2.0", features = ["axum_extras"] }
 utoipa-swagger-ui = { version = "6", features = ["axum"] }
-hyper = { version = "1.2.0", features = ["full"] }
+hyper = { version = "1.2.0", features = ["full"] }
+postgres = { version = "0.19.7"}
+sea-orm = {version = "0.12", features = ["runtime-tokio-rustls", "sqlx-postgres"]}

README.md

@@ -33,18 +33,18 @@ cargo run --bin scan -- -help
 ```
 
 ```shell
-Usage: scan.exe [OPTIONS] --repo <REPO>
+Usage: scan [OPTIONS] --repo <REPO>
 
 Options:
       --repo <REPO>                    Target repository
       --config <CONFIG>                Config path [default: gitleaks.toml]
       --threads <THREADS>              Maximum number of threads sensleak spawns [default: 10]
-      --chunk <BATCH_SIZE>             The number of git files processed in each batch [default: 10]
+      --chunk <CHUNK>                  The number of files processed in each batch [default: 10]
       --report <REPORT>                Path to write json leaks file
       --report-format <REPORT_FORMAT>  json, csv, sarif [default: json]
   -v, --verbose                        Show verbose output from scan
       --pretty                         Pretty print json if leaks are present
-      --commit <COMMIT>                sha of commit to scan or "latest" to scan the last commit of the repository
+      --commit <COMMIT>                sha of commit to scan
       --commits <COMMITS>              comma separated list of a commits to scan
       --commits-file <COMMITS_FILE>    file of new line separated list of a commits to scan
       --commit-since <COMMIT_SINCE>    Scan commits more recent than a specific date. Ex: '2006-01-02' or '2023-01-02T15:04:05-0700' format
@@ -57,6 +57,7 @@ Options:
       --repo-config                    Load config from target repo. Config file must be ".gitleaks.toml" or "gitleaks.toml"
       --debug                          log debug messages
       --disk <DISK>                    Clones repo(s) to disk
+      --to-db                          Output to database
   -h, --help                           Print help (see more with '--help')
   -V, --version                        Print version
 

src/entity/models.rs

@@ -2,6 +2,7 @@ use chrono::{DateTime, FixedOffset};
 use clap::Parser;
 use serde::{Deserialize, Serialize};
 use utoipa::{ToSchema};
+use sea_orm::{entity::prelude::*, ActiveValue};
 /// Represents the configuration for sensleaks tool.
 #[derive(Parser, Debug)]
 #[command(
@@ -98,6 +99,9 @@ pub struct Config {
     #[arg(long)]
     pub disk: Option<String>,
 
+    /// Output to database
+    #[arg(long)]
+    pub to_db: bool,
     // /// Start API
     // #[arg(long, default_value = "false")]
     // pub api: bool,
@@ -126,6 +130,7 @@ impl Default for Config {
             repo_config: false,
             debug: false,
             disk: None,
+            to_db: false,
             // api: false,
         }
     }
@@ -206,6 +211,31 @@ impl Default for Allowlist {
         Self::new()
     }
 }
+
+/// Sea-orm Entity
+#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel)]
+#[sea_orm(table_name = "leaks")]
+pub struct Model {
+    #[sea_orm(primary_key)]
+    pub id: i32,
+    pub line: String,
+    pub line_number: u32,
+    pub offender: String,
+    pub commit: String,
+    pub repo: String,
+    pub rule: String,
+    pub commit_message: String,
+    pub author: String,
+    pub email: String,
+    pub file: String,
+    pub date: String,
+}
+
+#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
+pub enum Relation {}
+
+impl ActiveModelBehavior for ActiveModel {}
+
 /// Represents an item in the scanned output.
 #[derive(Debug, Serialize, Deserialize, Clone)]
 pub struct Leak {
@@ -243,6 +273,25 @@ pub struct Leak {
     pub date: String,
 }
 
+impl Leak {
+    pub fn to_active_model(&self) -> ActiveModel {
+        ActiveModel {
+            line: ActiveValue::set(self.line.clone()),
+            line_number: ActiveValue::set(self.line_number),
+            offender: ActiveValue::set(self.offender.clone()),
+            commit: ActiveValue::set(self.commit.clone()),
+            repo: ActiveValue::set(self.repo.clone()),
+            rule: ActiveValue::set(self.rule.clone()),
+            commit_message: ActiveValue::set(self.commit_message.clone()),
+            author: ActiveValue::set(self.author.clone()),
+            email: ActiveValue::set(self.email.clone()),
+            file: ActiveValue::set(self.file.clone()),
+            date: ActiveValue::set(self.date.clone()),
+            ..Default::default()
+        }
+    }
+}
+
 /// The scan condition
 #[derive(Debug, Clone)]
 pub struct Scan {
@@ -356,3 +405,37 @@ pub struct CsvResult {
     /// The date of the commit.
     pub date: String,
 }
+
+/// Config to connect to the database
+#[derive(Debug, Default,Serialize, Deserialize)]
+pub struct ConnectDbConfig {
+    /// The host of the database
+    pub host: String,
+    /// The user of the database
+    pub user: String,
+    /// The password of the database
+    pub password: String,
+    /// The name of the database
+    pub dbname: String,
+    /// The port of the database
+    pub port: String,
+}
+
+impl ConnectDbConfig {
+    /// Translate the config to connection url
+    pub fn to_connection_url(&self) -> String {
+        format!(
+            "postgresql://{}:{}@{}:{}/{}",
+            self.user, self.password, self.host, self.port, self.dbname
+        )
+    }
+    pub fn new() -> Self {
+        ConnectDbConfig {
+            host: String::from(""),
+            user: String::from(""),
+            password: String::from(""),
+            dbname: String::from(""),
+            port: String::from(""),
+        }
+    }
+}

src/lib.rs

@@ -1,48 +1,43 @@
 mod errors;
- 
+
 mod utils {
     pub mod detect_utils;
     pub mod git_util;
-
 }
+
 pub mod entity{
     pub mod models;
 }
+
 pub mod service{
     pub mod detect_service;
     pub mod git_service;
+    pub mod db_service;
 }
 
 pub use entity::models;
- 
-
 pub use errors::*;
- 
- 
 pub use utils::detect_utils;
 pub use utils::git_util;
- 
 pub use git_util::*;
 pub use models::*;
 
 use axum::{routing, Router};
+
 use utoipa::{
      OpenApi,
 };
+
 use utoipa_swagger_ui::SwaggerUi;
 
- 
 mod routes{
     pub mod scan;
     pub mod rules;
 }
 pub use routes::scan::*;
 pub use routes::rules::*;
 
- use crate::routes::*;
-
- 
- 
+use crate::routes::*;
 
 pub async fn start() -> Result<(), Box<dyn std::error::Error>> {
     #[derive(OpenApi)]

src/main.rs

@@ -1,8 +1,9 @@
 use sensleak::service::detect_service::sensleaks;
 
 /// The entry of the project
-fn main() {
-    sensleaks();
+#[tokio::main]
+async fn main() {
+    sensleaks().await;
 }
 
 

src/routes/scan.rs

@@ -46,6 +46,9 @@ pub struct ConfigDto {
 
     /// Clones repo(s) to disk.
     pub disk: Option<String>,
+
+    /// Output to database
+    pub to_db: bool,
 }
 
 /// The return results of the scan.
@@ -94,8 +97,9 @@ pub async fn scan_repo(Json(json_config): Json<ConfigDto>) -> Json<ScanResponse>
     config.user = json_config.user;
     config.disk = json_config.disk;
     config.repo_config = json_config.repo_config.unwrap_or(false);
+    config.to_db = json_config.to_db;
 
-    match detect(config) {
+    match detect(config).await {
         Ok(results) => Json(ScanResponse {
             code: 200,
             leaks_number: Some(results.outputs.len()),

src/service/db_service.rs

@@ -0,0 +1,97 @@
+use crate::models::{ConnectDbConfig, Entity as Leaks, Leak};
+use chrono::Local;
+use sea_orm::*;
+use std::env;
+
+/// Sets up the database connection using the application's configuration settings.
+///
+/// This function attempts to establish a connection to the database using environment variables for the database configuration.
+/// It reads configuration values such as host, port, user, password, and database name from environment variables
+/// and uses them to construct the database URL.
+///
+/// # Returns
+///
+/// Returns a `Result<DatabaseConnection, DbErr>`:
+/// - `Ok(DatabaseConnection)` if the connection is successfully established.
+/// - `Err(DbErr)` if there is an error connecting to the database.
+pub async fn set_up_db() -> Result<DatabaseConnection, DbErr> {
+    let config = get_db_config();
+    let db_url = config.to_connection_url();                    
+    let db = Database::connect(&db_url).await?;
+    Ok(db)
+}
+
+/// Inserts a vector of `Leak` entities into the database and ensures that the `Leaks` table exists.
+///
+/// This function first checks if the `Leaks` table exists in the database and creates it if not.
+/// Then, it proceeds to insert the provided vector of `Leak` entities into the `Leaks` table.
+///
+/// # Arguments
+///
+/// * `_leaks` - A reference to a vector of `Leak` entities to be inserted into the database.
+///
+/// # Returns
+///
+/// Returns a `Result<(), DbErr>` indicating the outcome of the operation:
+/// - `Ok(())` if the insertion is successful and the `Leaks` table is either found or successfully created.
+/// - `Err(DbErr)` if there is an error during the table check/creation or insertion process.
+pub async fn insert_leaks(_leaks: &[Leak]) -> Result<(), DbErr> {
+    let db = match set_up_db().await {
+        Ok(db) => db,
+        Err(err) => panic!("{}", err),
+    };
+
+    // Check if the table Leaks exists and create it if not
+    let builder = db.get_database_backend();
+    let schema = Schema::new(builder);
+
+    let stmt = schema
+        .create_table_from_entity(Leaks)
+        .if_not_exists()
+        .to_owned();
+
+    let stmt = builder.build(&stmt);
+
+    db.execute(stmt).await?;
+      
+    println!(
+        "\x1b[34m[INFO]\x1b[0m[{}] Create Success ...",
+        Local::now().format("%Y-%m-%d %H:%M:%S"),
+    );
+
+    // Insert leaks
+    for leak in _leaks.iter() {
+        let active_model = leak.to_active_model();
+
+        let insert_result = Leaks::insert(active_model)
+            .exec(&db)
+            .await?;
+        println!("Inserted leak with result: {:?}", insert_result);
+    }
+
+    println!(
+        "\x1b[34m[INFO]\x1b[0m[{}] Insert Success ...",
+        Local::now().format("%Y-%m-%d %H:%M:%S"),
+    );
+
+    Ok(())
+}
+
+/// Retrieves database connection configuration from environment variables.
+///
+/// This function constructs a `ConnectDbConfig` struct with database connection details
+/// such as host, port, username, password, and database name, reading the values from
+/// environment variables. If an environment variable is not set, it defaults to a predefined value.
+///
+/// # Returns
+///
+/// Returns a `ConnectDbConfig` struct populated with the database connection details.
+fn get_db_config() -> ConnectDbConfig {
+    ConnectDbConfig { 
+        host: env::var("PG_HOST").unwrap_or("localhost".to_string()), 
+        port: env::var("PG_PORT").unwrap_or("5432".to_string()), 
+        user: env::var("PG_USER").unwrap_or("postgres".to_string()), 
+        password: env::var("PG_PASSWORD").unwrap_or("postgres".to_string()), 
+        dbname: env::var("PG_DBNAME").unwrap_or("postgres".to_string()) 
+    }
+}