Skip to content

Commit ec87f06

Browse files
okynosokynos
committed
Added hash update and new file insert on second scans
1 parent 8d3728f commit ec87f06

File tree

8 files changed

+336
-89
lines changed

8 files changed

+336
-89
lines changed

‎Cargo.toml‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ itertools = "0.13.0"
1414
hex = "0.4.3"
1515
futures = "0.3.21"
1616
sha3 = { version = "0.10.0", default-features = false }
17-
sha2 = { version = "0.10.8" }
1817
log = { version = "0.4.11", default-features = false }
1918
gethostname = { version = "0.5.0", default-features = false }
2019
uuid = { version = "1.0.0", default-features = false, features = ["v4"] }
@@ -24,8 +23,8 @@ tokio-util = { version = "0.7.8", default-features = false, features = ["codec"]
2423
serde_json = { version = "1.0.79", default-features = false }
2524
time = { version = "0.3.17", default-features = false }
2625
ctrlc = { version = "3.3.1", default-features = false, features = ["termination"] }
27-
log-panics = { version = "2.1.0", features = ["with-backtrace"]}
28-
rusqlite = { version = "0.32.1", features = ["bundled"]}
26+
log-panics = { version = "2.1.0", features = ["with-backtrace"] }
27+
rusqlite = { version = "0.32.1", features = ["bundled"] }
2928
walkdir = "2.5.0"
3029

3130
[dependencies.regex]
@@ -36,10 +35,12 @@ features = ["std"]
3635
[target.'cfg(windows)'.dependencies]
3736
windows-service = "0.7.0"
3837
zip = "2.1.3"
38+
sha2 = { version = "0.10.8" }
3939

4040
[target.'cfg(unix)'.dependencies]
4141
flate2 = "1.0.27"
4242
tar = "0.4.40"
43+
sha2 = { version = "0.10.8", features = ["asm"] }
4344

4445
[dev-dependencies]
4546
tokio-test = "*"

‎src/appconfig.rs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ pub struct AppConfig {
3232
pub events_destination: String,
3333
pub events_max_file_checksum: usize,
3434
pub events_max_file_size: usize,
35+
pub checksum_method: String,
3536
pub endpoint_type: String,
3637
pub endpoint_address: String,
3738
pub endpoint_user: String,
@@ -60,6 +61,7 @@ impl AppConfig {
6061
events_destination: self.events_destination.clone(),
6162
events_max_file_checksum: self.events_max_file_checksum,
6263
events_max_file_size: self.events_max_file_size,
64+
checksum_method: self.checksum_method.clone(),
6365
endpoint_type: self.endpoint_type.clone(),
6466
endpoint_address: self.endpoint_address.clone(),
6567
endpoint_user: self.endpoint_user.clone(),
@@ -128,6 +130,9 @@ impl AppConfig {
128130
None => 128
129131
};
130132

133+
// Temporal value
134+
let checksum_method = String::from("Partial");
135+
131136
// Manage null value on events->endpoint->insecure value
132137
let insecure = match yaml[0]["events"]["endpoint"]["insecure"].as_bool() {
133138
Some(value) => value,
@@ -273,6 +278,7 @@ impl AppConfig {
273278
events_destination,
274279
events_max_file_checksum,
275280
events_max_file_size,
281+
checksum_method,
276282
endpoint_type,
277283
endpoint_address,
278284
endpoint_user,
@@ -514,6 +520,8 @@ mod tests {
514520
events_destination: String::from(events_destination),
515521
events_max_file_checksum: 64,
516522
events_max_file_size: 128,
523+
checksum_method: String::from("Partial"),
524+
//////////////////////////////////////////////////////Include checksum method tests
517525
endpoint_type: String::from("Elastic"),
518526
endpoint_address: String::from("test"),
519527
endpoint_user: String::from("test"),

‎src/db.rs‎

Lines changed: 101 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,14 @@
1-
use rusqlite::Connection;
2-
use std::path::Path;
31
use crate::appconfig;
42
use crate::utils;
3+
use crate::dbfile::*;
4+
5+
use rusqlite::{Connection, Error};
6+
use rusqlite::Error::QueryReturnedNoRows;
7+
use std::path::Path;
58
use log::*;
6-
use std::fmt;
79

8-
pub const DBNAME: &str = "fim.db";
910

10-
pub struct DBFile {
11-
pub id: u64,
12-
pub timestamp: String,
13-
pub hash: String,
14-
pub path: String,
15-
pub size: u64
16-
}
11+
pub const DBNAME: &str = "fim.db";
1712

1813
pub struct DB {
1914
path: String
@@ -57,6 +52,35 @@ impl DB {
5752

5853
// ------------------------------------------------------------------------
5954

55+
pub fn exists(&self) -> bool {
56+
let mut config_folder = Path::new(&appconfig::get_config_path(utils::get_os()))
57+
.parent().unwrap().to_path_buf();
58+
config_folder.push(DBNAME);
59+
60+
config_folder.exists()
61+
}
62+
63+
// ------------------------------------------------------------------------
64+
65+
pub fn is_empty(&self) -> bool {
66+
let connection = self.open();
67+
let result = connection.query_row("SELECT * FROM files LIMIT 1", [], |_row| Ok(0));
68+
self.close(connection);
69+
match result {
70+
Ok(_v) => false,
71+
Err(e) => {
72+
if e == QueryReturnedNoRows {
73+
true
74+
} else {
75+
error!("Could not check if the database is empty, Error: {}", e);
76+
true
77+
}
78+
}
79+
}
80+
}
81+
82+
// ------------------------------------------------------------------------
83+
6084
pub fn create_table(&self) {
6185
let connection = self.open();
6286
let result = connection.execute(
@@ -92,11 +116,44 @@ impl DB {
92116

93117
// ------------------------------------------------------------------------
94118

95-
pub fn get_file(&self, path: String) -> DBFile {
119+
pub fn get_file_by_path(&self, path: String) -> Result<DBFile, DBFileError> {
96120
let connection = self.open();
97-
let data = connection.query_row(
121+
let result = connection.query_row(
98122
"SELECT * FROM files WHERE path = ?1 LIMIT 1",
99-
[path],
123+
[path.clone()],
124+
|row| Ok(DBFile {
125+
id: row.get(0).unwrap(),
126+
timestamp: row.get(1).unwrap(),
127+
hash: row.get(2).unwrap(),
128+
path: row.get(3).unwrap(),
129+
size: row.get(4).unwrap()
130+
})
131+
);
132+
133+
let data = match result {
134+
Ok(d) => Ok(d),
135+
Err(e) => {
136+
match e {
137+
Error::QueryReturnedNoRows => Err(DBFileError::not_found_error()),
138+
_ => {
139+
error!("Could not get file '{}' information in database, Error: {:?}", path, e);
140+
Err(DBFileError::from(e))
141+
}
142+
}
143+
}
144+
};
145+
146+
self.close(connection);
147+
data
148+
}
149+
150+
// ------------------------------------------------------------------------
151+
152+
pub fn get_file_by_id(&self, id: u64) -> DBFile {
153+
let connection = self.open();
154+
let data = connection.query_row(
155+
"SELECT * FROM files WHERE id = ?1 LIMIT 1",
156+
[id],
100157
|row| Ok(DBFile {
101158
id: row.get(0).unwrap(),
102159
timestamp: row.get(1).unwrap(),
@@ -112,6 +169,36 @@ impl DB {
112169

113170
// ------------------------------------------------------------------------
114171

172+
pub fn update_file(&self, dbfile: DBFile, timestamp: Option<String>, hash: Option<String>, size: Option<u64>) {
173+
let connection = self.open();
174+
175+
let timestamp_str = match timestamp {
176+
Some(t) => format!("timestamp = '{}'", t),
177+
None => String::new()
178+
};
179+
let hash_str = match hash {
180+
Some(h) => format!("hash = '{}'", h),
181+
None => String::new()
182+
};
183+
let size_str = match size {
184+
Some(s) => format!("size = '{}'", s),
185+
None => String::new()
186+
};
187+
188+
let query = format!("UPDATE files SET {}, {}, {} WHERE id = {}",
189+
timestamp_str, hash_str, size_str, dbfile.id);
190+
191+
let mut statement = connection.prepare(&query).unwrap();
192+
let result = statement.execute([]);
193+
match result {
194+
Ok(_v) => debug!("File '{}', updated with new information.", dbfile.path),
195+
Err(e) => error!("Cannot update file '{}' information, Error: {:?}", dbfile.path, e)
196+
}
197+
198+
}
199+
200+
// ------------------------------------------------------------------------
201+
115202
pub fn print(&self) {
116203
let connection = self.open();
117204
let mut query = connection.prepare(
@@ -130,18 +217,4 @@ impl DB {
130217
println!("{:?}", file.unwrap());
131218
}
132219
}
133-
}
134-
135-
// ----------------------------------------------------------------------------
136-
137-
impl fmt::Debug for DBFile {
138-
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result{
139-
f.debug_tuple("")
140-
.field(&self.id)
141-
.field(&self.timestamp)
142-
.field(&self.hash)
143-
.field(&self.path)
144-
.field(&self.size)
145-
.finish()
146-
}
147220
}

‎src/dbfile.rs‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
use crate::utils;
2+
use crate::hash;
3+
use crate::appconfig::*;
4+
5+
use sha2::{Digest, Sha256};
6+
use std::fmt;
7+
use std::path::Path;
8+
use rusqlite;
9+
10+
pub struct DBFileError {
11+
kind: String,
12+
message: String
13+
}
14+
15+
pub struct DBFile {
16+
pub id: u64,
17+
pub timestamp: String,
18+
pub hash: String,
19+
pub path: String,
20+
pub size: u64
21+
}
22+
23+
// ----------------------------------------------------------------------------
24+
25+
impl DBFileError {
26+
pub fn not_found_error() -> Self {
27+
DBFileError {
28+
kind: String::from("DBFileNotFoundError"),
29+
message: String::from("Could not find requested file in the database."),
30+
}
31+
}
32+
33+
// ------------------------------------------------------------------------
34+
35+
pub fn kind(&self) -> String {
36+
self.kind.clone()
37+
}
38+
}
39+
40+
// ----------------------------------------------------------------------------
41+
42+
impl From<rusqlite::Error> for DBFileError {
43+
fn from(error: rusqlite::Error) -> Self {
44+
DBFileError {
45+
kind: String::from("RusqliteError"),
46+
message: error.to_string()
47+
}
48+
}
49+
}
50+
51+
// ----------------------------------------------------------------------------
52+
53+
impl fmt::Debug for DBFileError {
54+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result{
55+
f.debug_tuple("")
56+
.field(&self.kind)
57+
.field(&self.message)
58+
.finish()
59+
}
60+
}
61+
62+
// ----------------------------------------------------------------------------
63+
64+
impl fmt::Debug for DBFile {
65+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result{
66+
f.debug_tuple("")
67+
.field(&self.id)
68+
.field(&self.timestamp)
69+
.field(&self.hash)
70+
.field(&self.path)
71+
.field(&self.size)
72+
.finish()
73+
}
74+
}
75+
76+
// ----------------------------------------------------------------------------
77+
78+
impl fmt::Display for DBFile {
79+
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
80+
write!(f, "Oh no, something bad went down")
81+
}
82+
}
83+
84+
// ----------------------------------------------------------------------------
85+
86+
impl DBFile {
87+
pub fn new(cfg: AppConfig, path: &str) -> Self {
88+
let size = Path::new(path).metadata().unwrap().len();
89+
let hash = match cfg.clone().checksum_method.as_str() {
90+
"Partial" => hash::get_partial_checksum(
91+
String::from(path),
92+
Sha256::new()
93+
),
94+
_ => hash::get_checksum(
95+
String::from(path),
96+
cfg.clone().events_max_file_checksum,
97+
Sha256::new())
98+
};
99+
100+
DBFile {
101+
id: 0,
102+
timestamp: utils::get_current_time_millis(),
103+
hash,
104+
path: String::from(path),
105+
size
106+
}
107+
}
108+
109+
// ------------------------------------------------------------------------
110+
111+
pub fn clone(&self) -> Self {
112+
DBFile {
113+
id: self.id,
114+
timestamp: self.timestamp.clone(),
115+
hash: self.hash.clone(),
116+
path: self.path.clone(),
117+
size: self.size
118+
}
119+
}
120+
121+
// ------------------------------------------------------------------------
122+
123+
pub fn get_disk_hash(&self, cfg: AppConfig) -> String {
124+
match cfg.clone().checksum_method.as_str() {
125+
"Partial" => hash::get_partial_checksum(
126+
String::from(&self.path),
127+
Sha256::new()
128+
),
129+
_ => hash::get_checksum(
130+
String::from(&self.path),
131+
cfg.clone().events_max_file_checksum,
132+
Sha256::new())
133+
}
134+
}
135+
}

0 commit comments

Comments
 (0)