From 9d218500338a2814979a9a40dd942254633e7858 Mon Sep 17 00:00:00 2001
From: FlyLoongZ <cnzy2004@foxmail.com>
Date: Fri, 27 Dec 2024 21:47:41 +0800
Subject: [PATCH] =?UTF-8?q?=E6=B7=BB=E5=8A=A0sqlite=E6=A0=BC=E5=BC=8Faudit?=
 =?UTF-8?q?=E6=96=87=E4=BB=B6=E8=BE=93=E5=87=BA?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/config/parser/mod.rs |  2 +-
 src/dns_mw_audit.rs      | 57 +++++++++++++++++++++++++++++++++++++---
 2 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/src/config/parser/mod.rs b/src/config/parser/mod.rs
index 572470fa..3a4dd6b0 100644
--- a/src/config/parser/mod.rs
+++ b/src/config/parser/mod.rs
@@ -150,7 +150,7 @@ pub fn parse_config(input: &str) -> IResult<&str, OneConfig> {
 
     fn parse_item<'a, T: NomParser>(
         keyword: &'static str,
-    ) -> impl FnMut(&'a str) -> IResult<&str, T> {
+    ) -> impl FnMut(&'a str) -> IResult<&'a str, T> {
         preceded(tuple((space0, tag_no_case(keyword), space1)), T::parse)
     }
 
diff --git a/src/dns_mw_audit.rs b/src/dns_mw_audit.rs
index 45e8aaca..dbbb3190 100644
--- a/src/dns_mw_audit.rs
+++ b/src/dns_mw_audit.rs
@@ -4,6 +4,8 @@ use std::path::Path;
 use std::time::Duration;
 use std::time::Instant;
 
+use rusqlite::{Connection, params};
+
 use chrono::prelude::*;
 use smallvec::SmallVec;
 use tokio::sync::mpsc::{self, Sender};
@@ -166,8 +168,6 @@ fn record_audit_to_file(
     audit_records: &[DnsAuditRecord],
 ) -> io::Result<()> {
     if matches!(audit_file.extension(), Some(ext) if ext == "csv") {
-        // write as csv
-
         if audit_file.peamble().is_none() {
             let mut writer = csv::Writer::from_writer(vec![]);
             writer.write_record([
@@ -211,8 +211,13 @@ fn record_audit_to_file(
                 format!("{:?}", audit.lookup_source).as_str(),
             ])?;
         }
+    } else if matches!(audit_file.extension(), Some(ext) if ext == "db") {
+        // 写入 SQLite 数据库
+        if let Err(e) = write_to_sqlite(audit_file.path(), audit_records) {
+            warn!("Failed to write audit records to SQLite database: {}", e);
+        }
     } else {
-        // write as nornmal log format.
+        // 写入普通日志格式
         for audit in audit_records {
             if writeln!(audit_file, "{}", audit).is_err() {
                 warn!("Write audit to file '{:?}' failed", audit_file.path());
@@ -223,6 +228,52 @@ fn record_audit_to_file(
     Ok(())
 }
 
+/// 将审计记录写入 SQLite 数据库
+fn write_to_sqlite(db_path: &Path, audit_records: &[DnsAuditRecord]) -> rusqlite::Result<()> {
+    // 打开或创建数据库连接
+    let conn = Connection::open(db_path)?;
+
+    // 创建表（如果不存在）
+    conn.execute(
+        "CREATE TABLE IF NOT EXISTS smartdns_audit (
+            id INTEGER PRIMARY KEY,
+            timestamp INTEGER NOT NULL,
+            client TEXT NOT NULL,
+            name TEXT NOT NULL,
+            type TEXT NOT NULL,
+            elapsed TEXT NOT NULL,
+            speed TEXT NOT NULL,
+            state TEXT NOT NULL,
+            result TEXT NOT NULL,
+            lookup_source TEXT NOT NULL
+        )",
+        [],
+    )?;
+
+    // 插入审计记录
+    for audit in audit_records {
+        conn.execute(
+            "INSERT INTO smartdns_audit (
+                id, timestamp, client, name, type, elapsed, speed, state, result, lookup_source
+            ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10)",
+            params![
+                audit.id,
+                audit.date.timestamp(),
+                audit.client,
+                audit.query.name().to_string(),
+                audit.query.query_type().to_string(),
+                format!("{:?}", audit.elapsed),
+                format!("{:?}", audit.speed),
+                if audit.result.is_ok() { "success" } else { "failed" },
+                audit.fmt_result(),
+                format!("{:?}", audit.lookup_source),
+            ],
+        )?;
+    }
+
+    Ok(())
+}
+
 #[cfg(test)]
 mod tests {