#[cfg(any(test, feature = "admin"))]
use super::admin;
#[cfg(any(test, feature = "admin"))]
use super::hash_password;
use super::{random_bytes_api_key, DatabaseInformation, MDBConfig};
use crate::crypto::{EncryptionOption, FileEncryption};
use crate::db::types::{FileMetadata, FileType};
use malwaredb_api::{
digest::HashType, GetUserInfoResponse, Label, Labels, SimilarSample, SimilarityHashType,
SourceInfo, Sources,
};
use malwaredb_types::KnownType;
use std::collections::HashMap;
use std::fmt::{Debug, Display, Formatter};
use std::path::PathBuf;
use std::str::FromStr;
use anyhow::{anyhow, bail, Context, Result};
use argon2::{Argon2, PasswordHash, PasswordVerifier};
#[cfg(any(test, feature = "admin"))]
use chrono::{DateTime, Local, Utc};
use fuzzyhash::FuzzyHash;
use humansize::{make_format, DECIMAL};
use malwaredb_lzjd::LZDict;
#[cfg(feature = "vt")]
use malwaredb_virustotal::filereport::ScanResultAttributes;
use rusqlite::params;
use rusqlite::{Batch, Connection};
use serde_json::json;
use tlsh_fixed::Tlsh;
use tracing::{debug, error};
const ARRAY_DELIMITER: &str = "|";
const SQLITE_SQL: &str = include_str!("malwaredb_sqlite.sql");
pub struct Sqlite {
conn: Connection,
file_path: PathBuf,
first_run: bool,
}
impl Sqlite {
pub fn new(file_name: &str) -> Result<Self> {
let file_path = std::path::Path::new(file_name);
let new_db = !file_path.exists();
if new_db {
let conn = Connection::open(file_name)?;
let mut batch = Batch::new(&conn, SQLITE_SQL);
while let Some(mut stmt) = batch.next()? {
stmt.execute([])?;
}
conn.execute("update mdbconfig set version = ?1", [&crate::MDB_VERSION])?;
}
Ok(Self {
conn: Connection::open(file_name)?,
file_path: file_name.into(),
first_run: new_db,
})
}
pub fn first_run(&self) -> bool {
self.first_run
}
pub fn enable_compression(&self) -> Result<()> {
if self.first_run {
self.conn
.execute("update mdbconfig set compress = true", ())?;
} else {
bail!("refusing to set compression when not on the first run");
}
Ok(())
}
#[cfg(feature = "vt")]
pub fn enable_vt_upload(&self) -> Result<()> {
if self.first_run {
self.conn
.execute("update mdbconfig set send_samples_to_vt = true", ())?;
} else {
bail!("refusing to set allow of sample uploading to VT when not on the first run");
}
Ok(())
}
#[cfg(feature = "vt")]
pub fn files_without_vt_records(&self, limit: i32) -> Result<Vec<String>> {
let mut statement = self
.conn.prepare("select file.sha256 from file where file.id not in (select fileid from vtdata) limit ?1")?;
let results = statement.query_map([limit], |row| {
let hash: String = row.get(0)?;
Ok(hash)
})?;
Ok(results.filter_map(|r| r.ok()).collect())
}
#[cfg(feature = "vt")]
pub fn store_vt_record(&self, results: &ScanResultAttributes) -> Result<()> {
let mut statement = self.conn.prepare("select id from file where sha256 = ?1")?;
let fid = statement.query_row([&results.sha256], |row| {
let fid: Option<i32> = row.get(0)?;
Ok(fid)
})?;
if let Some(id) = fid {
let report = serde_json::to_string(&results.last_analysis_results)?;
self.conn.execute(
"insert into vtdata(fileid, tstamp, hits, total, vtdetail) values(?1, ?2, ?3, ?4, ?5)",
(
id,
results.last_analysis_stats.malicious,
results.last_analysis_stats.av_count(),
report,
&results.last_analysis_date,
),
)?;
}
Ok(())
}
#[cfg(feature = "vt")]
pub fn get_vt_stats(&self) -> Result<super::VtStats> {
let mut statement = self
.conn
.prepare("select count(1) from vtdata where hits = 0")?;
let clean_records = statement.query_row([], |row| {
let clean: i32 = row.get(0)?;
Ok(clean as u32)
})?;
let mut statement = self
.conn
.prepare("select count(1) from vtdata where hits > 0")?;
let hits_records = statement.query_row([], |row| {
let hits: i32 = row.get(0)?;
Ok(hits as u32)
})?;
let mut statement = self.conn.prepare(
"select count(1) from file where file.id not in (select fileid from vtdata)",
)?;
let files_without_records = statement.query_row([], |row| {
let no_recs: i32 = row.get(0)?;
Ok(no_recs as u32)
})?;
Ok(super::VtStats {
clean_records,
hits_records,
files_without_records,
})
}
pub fn get_config(&self) -> Result<MDBConfig> {
let mut statement = self
.conn
.prepare("select name, compress, send_samples_to_vt, defaultKey from mdbconfig")?;
statement
.query_row([], |row| {
Ok(MDBConfig {
name: row.get(0)?,
compression: row.get(1)?,
send_samples_to_vt: row.get(2)?,
default_key: row.get(3)?,
})
})
.map_err(anyhow::Error::new)
}
pub fn authenticate(&self, uname: &str, password: &str) -> Result<String> {
let mut statement = self
.conn
.prepare("select password from person where uname = ?1")?;
let results = statement.query_row([uname], |row| {
let db_password: Option<String> = row.get(0)?;
Ok(db_password)
})?;
if let Some(password_hash) = results {
let password_hashed = PasswordHash::new(&password_hash)?;
Argon2::default().verify_password(password.as_ref(), &password_hashed)?;
} else {
bail!("Password not set");
}
let mut statement = self
.conn
.prepare("select apikey from person where uname = ?1")?;
let result = statement.query_row([uname], |row| {
let apikey: Option<String> = row.get(0)?;
if apikey.is_none() {
return Ok(None);
}
let key: String = apikey.unwrap();
Ok(Some(key))
})?;
if let Some(apikey) = result {
return Ok(apikey);
}
let apikey = random_bytes_api_key();
self.conn.execute(
"update person set apikey = ?1 where uname = ?2",
(&apikey, &uname),
)?;
Ok(apikey)
}
pub fn get_uid(&self, apikey: &str) -> Result<i32> {
let mut statement = self
.conn
.prepare("select id from person where apikey = ?1")?;
let uid = statement.query_row([apikey], |row| {
let uid: Option<i32> = row.get(0)?;
Ok(uid)
})?;
uid.ok_or(anyhow!("unable to get user ID"))
}
pub fn db_info(&self) -> Result<DatabaseInformation> {
let size = self.file_path.metadata()?.len();
let size = humansize::SizeFormatter::new(size, humansize::BINARY).to_string();
let mut statement = self.conn.prepare("SELECT sqlite_version();")?;
let version = statement.query_row([], |row| {
let ver: String = row.get(0)?;
Ok(ver)
})?;
let mut statement = self.conn.prepare("SELECT count(1) from person")?;
let num_users = statement.query_row([], |row| {
let users: u32 = row.get(0)?;
Ok(users)
})?;
let mut statement = self.conn.prepare("SELECT count(1) from file")?;
let num_files = statement.query_row([], |row| {
let files: u64 = row.get(0)?;
Ok(files)
})?;
let mut statement = self.conn.prepare("SELECT count(1) from grp")?;
let num_groups = statement.query_row([], |row| {
let groups: u32 = row.get(0)?;
Ok(groups)
})?;
let mut statement = self.conn.prepare("SELECT count(1) from source")?;
let num_sources = statement.query_row([], |row| {
let sources: u32 = row.get(0)?;
Ok(sources)
})?;
Ok(DatabaseInformation {
version: format!("SQLite {version}"),
size,
num_files,
num_users,
num_groups,
num_sources,
})
}
pub fn get_user_info(&self, uid: i32) -> Result<GetUserInfoResponse> {
let mut statement = self
.conn
.prepare("select uname from person where id = ?1")?;
let username = statement.query_row([uid], |row| {
let uname: Option<String> = row.get(0)?;
Ok(uname)
})?;
let username = username.ok_or(anyhow!("unable to get username"))?;
let mut is_admin = false;
let mut statement = self.conn.prepare("select grp.name, grp.id, grp.parent from grp, usergroup where grp.id = usergroup.gid and usergroup.pid = ?1")?;
let rows = statement.query_map([uid], |row| {
let name: String = row.get(0)?;
let id: i32 = row.get(1)?;
let parent_id: Option<i32> = row.get(2)?;
if id == 0 {
is_admin = true;
}
if let Some(parent) = parent_id {
if parent == 0 {
is_admin = true;
}
}
Ok(name)
})?;
let mut groups = vec![];
for row in rows {
groups.push(row?);
}
let mut statement = self.conn.prepare("select source.name from source, usergroup, groupsource where source.id = groupsource.sourceid and groupsource.gid = usergroup.gid and usergroup.pid = ?1")?;
let rows = statement.query_map([uid], |row| {
let name: String = row.get(0)?;
Ok(name)
})?;
let mut sources = vec![];
for row in rows {
sources.push(row?);
}
Ok(GetUserInfoResponse {
id: uid,
username,
groups,
sources,
is_admin,
})
}
pub fn get_user_sources(&self, uid: i32) -> Result<Sources> {
let mut statement = self.conn.prepare("select source.id, source.name, source.description, source.url, source.firstacquisition from source, usergroup, groupsource where source.id = groupsource.sourceid and groupsource.gid = usergroup.gid and usergroup.pid = ?1")?;
let mut sources = vec![];
let rows = statement.query_map([uid], |row| {
let id: i32 = row.get(0)?;
let name: String = row.get(1)?;
let description: Option<String> = row.get(2)?;
let url: Option<String> = row.get(3)?;
let date: String = row.get(4).expect("failed to get date");
let date = chrono::DateTime::parse_from_rfc3339(&date)
.expect("failed to get source date in in Sqlite::list_groups");
let first_acquisition = date.with_timezone(&chrono::Utc);
Ok(SourceInfo {
id: id as u32,
name,
description,
url,
first_acquisition,
})
})?;
for row in rows {
sources.push(row?);
}
Ok(Sources {
sources,
message: None,
})
}
pub fn reset_own_api_key(&self, uid: i32) -> Result<()> {
self.conn.execute(
"update person set apikey = NULL where id = ?1",
params![uid],
)?;
Ok(())
}
pub fn get_known_data_types(&self) -> Result<Vec<FileType>> {
let mut statement = self
.conn
.prepare("select id, name, description, magic, executable from filetype")?;
let rows = statement.query_map([], |row| {
let name: String = row.get(1)?;
let magic = {
let magic: String = row.get(3)?;
let mut magic_numbers = vec![];
if magic.contains(ARRAY_DELIMITER) {
for magic_hex in magic.split(ARRAY_DELIMITER) {
magic_numbers.push(
hex::decode(magic_hex)
.context(format!("could not decode hex from datatype {name}"))
.unwrap(),
);
}
} else {
magic_numbers.push(
hex::decode(magic)
.context(format!("could not decode hex from datatype {name}"))
.unwrap(),
);
}
magic_numbers
};
let exec: i32 = row.get(4)?;
Ok(FileType {
id: row.get(0)?,
name,
description: row.get(2)?,
magic,
executable: exec > 0,
})
})?;
let mut file_types = vec![];
for row in rows {
file_types.push(row?);
}
Ok(file_types)
}
pub fn get_labels(&self) -> Result<Labels> {
let mut statement = self.conn.prepare("select label.id, label.name, parent.name from label left outer join label parent on (parent.id = label.parent)")?;
let mut labels = Vec::new();
let rows = statement.query_map([], |row| {
let id: i64 = row.get(0)?;
let name: String = row.get(1)?;
let parent: Option<String> = row.get(2)?;
Ok((id, name, parent))
})?;
for (id, name, parent) in rows.flatten() {
labels.push(Label {
id: id as u64,
name,
parent,
});
}
Ok(Labels(labels))
}
pub fn get_type_id_for_bytes(&self, data: &[u8]) -> Result<i32> {
let db_file_types = match self.get_known_data_types() {
Ok(t) => t,
Err(e) => {
return Err(e);
}
};
for db_file_type in db_file_types {
for magic in db_file_type.magic {
if data.starts_with(&magic) {
return Ok(db_file_type.id);
}
}
}
bail!("File type not found")
}
pub fn allowed_user_source(&self, uid: i32, sid: i32) -> Result<bool> {
let mut statement = self.conn.prepare(
"select usergroup.gid from usergroup, groupsource where usergroup.gid = groupsource.gid and usergroup.pid = ?1 and groupsource.sourceid = ?2")?;
let results = statement.query_map([uid, sid], |row| {
let gid: u32 = row.get(0)?;
Ok(gid)
})?;
let allowed = results.flatten().next().is_some();
Ok(allowed)
}
pub fn user_is_admin(&self, uid: i32) -> Result<bool> {
let mut statement = self.conn.prepare(
"select usergroup.gid from usergroup join grp on (usergroup.gid = grp.id) where usergroup.pid = ?1 and (usergroup.gid = 0 or grp.parent = 0)")?;
let results = statement.query_map([uid], |row| {
let gid: u32 = row.get(0)?;
Ok(gid)
})?;
let admin = results.flatten().next().is_some();
Ok(admin)
}
pub fn add_file(
&self,
meta: &FileMetadata,
known_type: KnownType<'_>,
uid: i32,
sid: i32,
ftype: i32,
parent: Option<i64>,
) -> Result<bool> {
if !self.allowed_user_source(uid, sid)? {
bail!("uid {uid} not allowed to upload to sid {sid}");
}
let creation_date = known_type.created().map(|d| d.to_rfc3339());
let mut statement = self
.conn
.prepare("select count(1) from file where sha512 = ?1")?;
let result = statement.query_map([&meta.sha512], |row| {
let counts: i32 = row.get(0)?;
Ok(counts)
})?;
let exists = result.flatten().next().unwrap();
let (fid, new_file) = if exists > 0 {
let mut statement = self.conn.prepare("select id from file where sha512 = ?1")?;
let result = statement.query_map([&meta.sha512], |row| {
let fid: i32 = row.get(0)?;
Ok(fid)
})?;
let fid = result.flatten().next().unwrap();
(fid, false)
} else {
self.conn.execute("insert into file(sha1, sha256, sha384, sha512, md5, lzjd, ssdeep, sdhash, tlsh, humanhash, filecommand, createddate, filetypeid, size, entropy)\
values(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15)", (&meta.sha1, &meta.sha256, &meta.sha384, &meta.sha512, &meta.md5, &meta.lzjd, &meta.ssdeep, &meta.sdhash, &meta.tlsh, &meta.humanhash, &meta.file_command, &creation_date, &ftype, &meta.size, &meta.entropy))?;
let mut statement = self.conn.prepare("select id from file where sha512 = ?1")?;
let result = statement.query_map([&meta.sha512], |row| {
let fid: i32 = row.get(0)?;
Ok(fid)
})?;
let fid = result.flatten().next().unwrap();
if let Some(parent_id) = parent {
let fid = fid as i64;
let mut statement = self
.conn
.prepare("update file set parent = ?1 where id = ?2")?;
if let Err(e) = statement.execute([&parent_id, &fid]) {
error!("SQLite: failed to set parent id for a file: {e}");
}
}
if known_type.is_doc() {
if let Some(doc) = known_type.clone().doc() {
self.conn.execute("insert into pdf(fileid, title, author, pages, javascript, forms) values(?1, ?2, ?3, ?4, ?5, ?6)",
(&fid, &doc.title(), &doc.author(), &doc.pages(), &doc.has_javascript(), &doc.has_form()))?;
}
} else if known_type.is_exec() {
if let Some(exec) = known_type.clone().exec() {
self.conn.execute(
"insert into executable(fileid, sections) values(?1, ?2)",
(&fid, &exec.num_sections()),
)?;
}
}
if let Some(children) = known_type.children() {
for child in children {
let child_meta = FileMetadata::new(child.contents(), None);
if let Ok(type_id) = self.get_type_id_for_bytes(child.contents()) {
if let Err(e) =
self.add_file(&child_meta, child, uid, sid, type_id, Some(fid as i64))
{
error!("Sqlite: failed to insert record for child file: {e}");
}
}
}
}
(fid, true)
};
let mut statement = self
.conn
.prepare("select count(*) from filesource where fileid = ?1 and sourceid = ?2")?;
let result = statement.query_map([&fid, &sid], |row| {
let count: i32 = row.get(0)?;
Ok(count)
})?;
let count = result.flatten().next().unwrap();
if count == 0 {
self.conn.execute(
"insert into filesource(fileid, sourceid, userid) values(?1, ?2, ?3)",
(&fid, &sid, &uid),
)?;
}
if let Some(fname) = &meta.name {
let mut statement = self
.conn
.prepare("select filename from filesource where fileid = ?1 and sourceid = ?2")?;
let result = statement.query_map([&fid, &sid], |row| {
let fnames: String = row.get(0).unwrap_or_default();
Ok(fnames)
})?;
let mut file_names = result.flatten().next().unwrap();
let mut changed = false;
if file_names.is_empty() {
file_names.clone_from(fname);
changed = true;
} else if !file_names.contains(fname) {
file_names = format!("{file_names}{ARRAY_DELIMITER}{fname}");
changed = true;
}
if changed {
self.conn.execute(
"update filesource set filename = ?1 where fileid = ?2 and sourceid = ?3",
(&file_names, &fid, &sid),
)?;
}
}
Ok(new_file)
}
pub fn retrieve_sample(&self, uid: i32, hash: &HashType) -> Result<String> {
let mut statement = self.conn.prepare(&format!("select distinct sha256 from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid)\
join usergroup on (groupsource.gid = usergroup.gid)\
where file.{} = ?1 and usergroup.pid = ?2", hash.name()))?;
statement
.query_row(params![hash.the_hash(), uid], |row| {
let sha256: String = row.get(0)?;
Ok(sha256)
})
.map_err(|e| anyhow!(e.to_string()))
}
pub fn get_sample_report(&self, uid: i32, hash: &HashType) -> Result<malwaredb_api::Report> {
let mut statement = if cfg!(feature = "vt") {
self.conn.prepare(&format!("select md5, sha1, sha256, sha384, sha512, lzjd, tlsh, ssdeep, sdhash, humanhash, filecommand, size, entropy, vtdata.hits, vtdata.total, vtdata.vtdetail, vtdata.tstamp from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
left outer join vtdata on (vtdata.fileid = file.id) \
where file.{} = ?1 and usergroup.pid = ?2 \
order by vtdata.tstamp desc limit 1", hash.name()))?
} else {
self.conn.prepare(&format!("select md5, sha1, sha256, sha384, sha512, lzjd, tlsh, ssdeep, sdhash, humanhash, filecommand, size, entropy from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where file.{} = ?1 and usergroup.pid = ?2", hash.name()))?
};
statement
.query_row(params![&hash.the_hash(), &uid], |row| {
let bytes: i32 = row.get(11)?;
let formatter = make_format(DECIMAL);
let vt = if cfg!(feature = "vt") {
let hits: Option<u32> = row.get(13)?;
let total: Option<u32> = row.get(14)?;
let detail: Option<String> = row.get(15)?;
Some(malwaredb_api::VirusTotalSummary {
hits: hits.unwrap_or_default(),
total: total.unwrap_or_default(),
detail: detail.map(|d| json!(d)),
last_analysis_date: row.get(16)?,
})
} else {
None
};
Ok(malwaredb_api::Report {
md5: row.get(0)?,
sha1: row.get(1)?,
sha256: row.get(2)?,
sha384: row.get(3)?,
sha512: row.get(4)?,
lzjd: row.get(5)?,
tlsh: row.get(6)?,
ssdeep: row.get(7)?,
sdhash: row.get(8)?,
humanhash: row.get(9)?,
filecommand: row.get(10)?,
bytes: bytes as u32,
size: formatter(bytes as u32),
entropy: row.get(12)?,
vt,
})
})
.map_err(|e| anyhow!(e.to_string()))
}
pub fn find_similar_samples(
&self,
uid: i32,
sim: &[(SimilarityHashType, String)],
) -> Result<Vec<SimilarSample>> {
let mut results = HashMap::<String, Vec<(SimilarityHashType, f32)>>::new();
for (algo, hash_value) in sim.iter().map(|(sim, val)| (*sim, val)) {
match algo {
SimilarityHashType::LZJD => {
let lzjd_hash = match LZDict::from_base64_string(hash_value) {
Ok(l) => l,
Err(e) => {
debug!("Error getting LZJD from web query: {e:?}");
continue;
}
};
let mut statement = self.conn.prepare(
"select sha256, lzjd \
from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where usergroup.pid = ?1",
)?;
let rows = statement.query_map([uid], |row| {
let sha256: String = row.get(0)?;
let lzjd: String = row.get(1)?;
Ok((sha256, lzjd))
})?;
for row in rows.into_iter() {
let (sha256, lzjd) = row?;
let lzjd = match LZDict::from_base64_string(&lzjd) {
Ok(l) => l,
Err(e) => {
debug!("Error getting LZJD from db query: {e:?}");
continue;
}
};
let similarity = lzjd.similarity(&lzjd_hash) as f32;
if let Some(already) = results.get_mut(&sha256) {
already.push((algo, similarity));
} else {
results.insert(sha256, vec![(algo, similarity)]);
}
}
}
SimilarityHashType::ImportHash | SimilarityHashType::PEHash => {
let mut statement = if algo == SimilarityHashType::ImportHash {
self.conn.prepare(
"select file.sha256 \
from file join executable on (file.id = executable.file_id) \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where executable.importhash = ?1 and usergroup.pid = ?2",
)?
} else {
self.conn.prepare(
"select file.sha256 \
from file join executable on (file.id = executable.file_id) \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where executable.pehash = ?1 and usergroup.pid = ?2",
)?
};
let rows = statement.query_map(params![&hash_value, &uid], |row| {
let sha256: String = row.get(0)?;
Ok(sha256)
})?;
for row in rows.into_iter() {
let sha256 = row?;
if let Some(already) = results.get_mut(&sha256) {
already.push((algo, 100.0));
} else {
results.insert(sha256, vec![(algo, 100.0)]);
}
}
}
SimilarityHashType::FuzzyImportHash => {
let mut statement = self.conn.prepare(
"select file.sha256, executable.importhashfuzzy \
from file join executable on (file.id = executable.file_id) \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where usergroup.pid = ?1",
)?;
let rows = statement.query_map([&uid], |row| {
let sha256: String = row.get(0)?;
let fuzzy_imp_hash: String = row.get(1)?;
Ok((sha256, fuzzy_imp_hash))
})?;
for row in rows.into_iter() {
let (sha256, fuzzy_imp_hash) = row?;
if let Ok(similarity) = FuzzyHash::compare(hash_value, fuzzy_imp_hash) {
if let Some(already) = results.get_mut(&sha256) {
already.push((algo, similarity as f32));
} else {
results.insert(sha256, vec![(algo, similarity as f32)]);
}
}
}
}
SimilarityHashType::SSDeep => {
let mut statement = self.conn.prepare(
"select sha256, ssdeep from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where usergroup.pid = ?1",
)?;
let rows = statement.query_map([&uid], |row| {
let sha256: String = row.get(0)?;
let ssdeep: String = row.get(1)?;
Ok((sha256, ssdeep))
})?;
for row in rows.into_iter() {
let (sha256, ssdeep_db_hash) = row?;
if let Ok(similarity) = FuzzyHash::compare(hash_value, ssdeep_db_hash) {
if let Some(already) = results.get_mut(&sha256) {
already.push((algo, similarity as f32));
} else {
results.insert(sha256, vec![(algo, similarity as f32)]);
}
}
}
}
SimilarityHashType::TLSH => {
let tlsh_hash = match Tlsh::from_str(hash_value) {
Ok(t) => t,
Err(e) => {
debug!("Error getting Tlsh from web query: {e}");
continue;
}
};
let mut statement = self.conn.prepare(
"select sha256, tlsh from file \
join filesource on (file.id = filesource.fileid) \
join groupsource on (groupsource.sourceid = filesource.sourceid) \
join usergroup on (groupsource.gid = usergroup.gid) \
where usergroup.pid = ?1",
)?;
let rows = statement.query_map([&uid], |row| {
let sha256: String = row.get(0)?;
let tlsh: String = row.get(1)?;
Ok((sha256, tlsh))
})?;
for row in rows.into_iter() {
let (sha256, tlsh) = row?;
let tlsh = match Tlsh::from_str(&tlsh) {
Ok(t) => t,
Err(e) => {
debug!("Error getting Tlsh from db query: {e}");
continue;
}
};
let similarity = tlsh.diff(&tlsh_hash, true) as f32;
if let Some(already) = results.get_mut(&sha256) {
already.push((algo, similarity));
} else {
results.insert(sha256, vec![(algo, similarity)]);
}
}
}
SimilarityHashType::SDHash => {
error!("SDHash is not yet implemented in Rust.");
continue;
}
_ => {
debug!("SQLite similarity check got unknown hash {algo}");
continue;
}
}
}
Ok(results
.into_iter()
.map(|(sha256, algorithms)| SimilarSample { sha256, algorithms })
.collect())
}
pub(crate) fn get_encryption_keys(&self) -> Result<HashMap<u32, FileEncryption>> {
let mut statement = self
.conn
.prepare("select id, name, bytes from encryptionkey")?;
let mut keys = HashMap::new();
let results = statement.query_map([], |row| {
let id: u32 = row.get(0)?;
let name: String = row.get(1)?;
let key: String = row.get(2)?;
Ok((id, name, key))
})?;
for result in results {
let (id, name, key) = result?;
let bytes = hex::decode(key)?;
let key = FileEncryption::new(EncryptionOption::try_from(name.as_str())?, bytes)?;
keys.insert(id, key);
}
Ok(keys)
}
pub(crate) fn get_file_encryption_key_id(
&self,
hash: &str,
) -> Result<(Option<u32>, Option<Vec<u8>>)> {
let mut statement = self
.conn
.prepare("select key, nonce from file where sha256 = ?1")?;
let (id, nonce) = statement.query_row([hash], |row| {
let id: Option<u32> = row.get(0)?;
let nonce: Option<String> = row.get(1)?;
Ok((id, nonce))
})?;
let nonce = if let Some(nonce) = nonce {
Some(hex::decode(nonce)?)
} else {
None
};
Ok((id, nonce))
}
pub(crate) fn set_file_nonce(&self, hash: &str, nonce: &Option<Vec<u8>>) -> Result<()> {
let nonce = nonce.as_ref().map(hex::encode);
self.conn.execute(
"update file set nonce = ?2 where sha256 = ?1",
(&hash, &nonce),
)?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn add_file_encryption_key(&self, key: &FileEncryption) -> Result<u32> {
let bytes = hex::encode(key.key());
self.conn.execute(
"insert into encryptionkey(name, bytes) values(?1, ?2)",
(&key.name(), &bytes),
)?;
let mut statement = self
.conn
.prepare("select id from encryptionkey where bytes = ?1")?;
let key_id = statement.query_row([bytes], |row| {
let id: u32 = row.get(0)?;
Ok(id)
})?;
self.conn
.execute("update mdbconfig set defaultKey = ?1 ", params![key_id])?;
Ok(key_id)
}
#[cfg(any(test, feature = "admin"))]
pub fn get_encryption_key_names_ids(&self) -> Result<Vec<(u32, EncryptionOption)>> {
let mut statement = self.conn.prepare("select id, name from encryptionkey")?;
let mut keys = vec![];
let results = statement.query_map([], |row| {
let id: u32 = row.get(0)?;
let name: String = row.get(1)?;
Ok((id, name))
})?;
for result in results {
let (id, name) = result?;
let key_option: EncryptionOption = name.as_str().try_into()?;
keys.push((id, key_option));
}
Ok(keys)
}
#[cfg(any(test, feature = "admin"))]
pub fn create_user(
&self,
uname: &str,
fname: &str,
lname: &str,
email: &str,
password: Option<String>,
organisation: Option<String>,
) -> Result<u64> {
let mut statement = self
.conn
.prepare("select count() from person where uname = ?1")?;
let results = statement.query_map([uname], |row| {
let count: u32 = row.get(0)?;
Ok(count)
})?;
if let Some(count) = results.flatten().next() {
if count != 0 {
bail!("User already exists");
}
}
let now: DateTime<Utc> = Utc::now();
match password {
None => {
self.conn
.execute("insert into person(email, uname, firstname, lastname, organisation, created) values (?1, ?2, ?3, ?4, ?5, ?6);", (&email, &uname, &fname, &lname, &organisation, &now.to_rfc3339()))?;
}
Some(pass) => {
let password = hash_password(&pass)?;
self.conn
.execute("insert into person(email, uname, firstname, lastname, organisation, password, created) values (?1, ?2, ?3, ?4, ?5, ?6, ?7);", (&email, &uname, &fname, &lname, &organisation, &password, &now.to_rfc3339()))?;
}
};
let mut statement = self
.conn
.prepare("select id from person where uname = ?1")?;
let results = statement.query_map([uname], |row| {
let uid: u64 = row.get(0)?;
Ok(uid)
})?;
if let Some(uid) = results.flatten().next() {
return Ok(uid);
}
bail!("User not created")
}
#[cfg(any(test, feature = "admin"))]
pub fn reset_api_keys(&self) -> Result<u64> {
let reset = self.conn.execute("update person set apikey = NULL", ())?;
Ok(reset as u64)
}
#[cfg(any(test, feature = "admin"))]
pub fn set_password(&self, uname: &str, password: &str) -> Result<()> {
let password = hash_password(password)?;
debug_assert_eq!(
self.conn
.execute(
"update person set password = ?1 where uname = ?2",
(&password, &uname),
)
.unwrap(),
1
);
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn list_users(&self) -> Result<Vec<admin::User>> {
let mut users = Vec::new();
let mut statement = self.conn.prepare("select id,email,uname,firstname,lastname,password is not null and length(password)>0,apikey is not null and length(apikey)>0,organisation,phone from person")?;
for result in statement.query_map([], |row| {
Ok(admin::User {
id: row.get(0)?,
email: row.get(1)?,
uname: row.get(2)?,
fname: row.get(3)?,
lname: row.get(4)?,
has_password: row.get(5)?,
has_api_key: row.get(6)?,
org: row.get(7)?,
phone: row.get(8)?,
})
})? {
if let Ok(user) = result {
users.push(user);
} else {
bail!("Failed to fetch user data");
}
}
Ok(users)
}
#[cfg(any(test, feature = "admin"))]
pub fn group_id_from_name(&self, name: &str) -> Result<i32> {
let mut statement = self.conn.prepare("select id from grp where name = ?1")?;
let id = statement.query_row(params![name], |row| {
let id: i32 = row.get(0)?;
Ok(id)
})?;
Ok(id)
}
#[cfg(any(test, feature = "admin"))]
pub fn edit_group(&self, gid: i32, name: &str, desc: &str, parent: Option<i32>) -> Result<()> {
self.conn.execute(
"update grp set name = ?1, description = ?2, parent = ?3 where id = ?4",
params![&name, &desc, &parent, &gid],
)?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn list_groups(&self) -> Result<Vec<admin::Group>> {
let mut groups = Vec::new();
let mut statement = self.conn.prepare("select grp.id, grp.name, grp.description, parent.name from grp left join grp parent on grp.parent = parent.id")?;
for result in statement.query_map([], |row| {
let id: i32 = row.get(0)?;
let members = {
let mut statement_members = self.conn.prepare("select person.id, person.uname, person.email, person.firstname, person.lastname, person.organisation, person.phone, person.password is not null and length(person.password)>0,person.apikey is not null and length(person.apikey)>0 from person, usergroup where person.id = usergroup.pid and usergroup.gid = ?1").expect("failed to get user names for group membership");
let mut members_list = vec![];
for member in statement_members.query_map([id], |member_row| {
Ok(admin::User {
id: member_row.get(0).expect("failed to get id"),
uname: member_row.get(1).expect("failed to get uname"),
email: member_row.get(2).expect("failed to get email"),
fname: member_row.get(3).expect("failed to get first name"),
lname: member_row.get(4).expect("failed to get last name"),
org: member_row.get(5).expect("failed to get org"),
phone: member_row.get(6).expect("failed to get phone"),
has_password: member_row.get(7).expect("failed to get password info"),
has_api_key: member_row.get(8).expect("failed to get api key info"),
})
})? {
members_list.push(member.expect("failed to get string form of uname"));
}
members_list
};
let sources = {
let mut sources_list = vec![];
let mut statement_sources = self.conn.prepare("select source.id, source.name, source.description, source.url, source.firstacquisition, parent_source.name from source left join source as parent_source on (source.parent = parent_source.id), groupsource where source.id = groupsource.sourceid and groupsource.gid = ?1")?;
for source in statement_sources.query_map([id], |source_row| {
let date: String = source_row.get(4).expect("failed to get date");
let date = chrono::DateTime::parse_from_rfc3339(&date).expect("failed to get source date in in Sqlite::list_groups");
let date = date.with_timezone(&Local);
let id = source_row.get(0).expect("failed to get id");
let mut counts_statement = self.conn.prepare("select (select count(1) from filesource where sourceid = ?1), (select count(1) from groupsource where sourceid = ?2)").unwrap();
let (files, groups) = counts_statement.query_row([id, id], |counts_row| {
let files: u32 = counts_row.get(0).unwrap();
let groups: u32 = counts_row.get(1).unwrap();
Ok((files, groups))
})?;
Ok(admin::Source {
id,
name: source_row.get(1).expect("failed to get name"),
description: source_row.get(2).expect("failed to get description"),
url: source_row.get(3).expect("failed to get url"),
date,
files,
groups,
parent: source_row.get(5).expect("failed to get source parent"),
})
})? {
sources_list.push(source.expect("unable to get source linked to a group"));
}
sources_list
};
let mut files_statement = self.conn.prepare("select count(filesource.fileid) from filesource join groupsource on groupsource.gid = filesource.sourceid where groupsource.gid = ?1")?;
let files = files_statement.query_row([id], |row| {
let count: u32 = row.get(0).expect("failed to get files count for group");
Ok(count)
})?;
Ok(admin::Group {
id,
name: row.get(1)?,
description: row.get(2)?,
parent: row.get(3)?,
members,
sources,
files,
})
})? {
match result {
Ok(group) => groups.push(group),
Err(e) => bail!("Failed to fetch group data: {e}"),
}
}
Ok(groups)
}
#[cfg(any(test, feature = "admin"))]
pub fn add_user_to_group(&self, uid: i32, gid: i32) -> Result<()> {
self.conn
.execute("insert into usergroup(pid, gid) values(?1, ?2)", (uid, gid))?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn add_group_to_source(&self, gid: i32, sid: i32) -> Result<()> {
self.conn.execute(
"insert into groupsource(gid, sourceid) values(?1, ?2)",
(gid, sid),
)?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn create_group(&self, name: &str, description: &str, parent: Option<i32>) -> Result<i32> {
let result = if let Some(pid) = parent {
self.conn.execute(
"insert into grp(name, description, parent) values (?1, ?2, ?3);",
params![&name, &description, &pid],
)
} else {
self.conn.execute(
"insert into grp(name, description) values (?1, ?2);",
[&name, &description],
)
}?;
if result != 1 {
bail!("failed to create group {name}");
}
let mut statement = self.conn.prepare("select id from grp where name = ?1")?;
let gid = statement
.query_map([name], |row| {
let gid: i32 = row.get(0)?;
Ok(gid)
})?
.flatten()
.next();
gid.ok_or(anyhow!("unable to get group ID"))
}
#[cfg(any(test, feature = "admin"))]
pub fn list_sources(&self) -> Result<Vec<admin::Source>> {
let mut sources = Vec::new();
let mut statement = self
.conn
.prepare("select source.id, source.name, source.description, source.url, source.firstacquisition, parent_source.name from source left join source as parent_source on (source.parent = parent_source.id)")?;
for result in statement.query_map([], |row| {
let id = row.get(0)?;
let mut counts_statement = self.conn.prepare("select (select count(1) from filesource where sourceid = ?1), (select count(1) from groupsource where sourceid = ?1)")?;
let (files, groups) = counts_statement.query_row([id], |counts_row| {
let files: u32 = counts_row.get(0)?;
let groups: u32 = counts_row.get(1)?;
Ok((files, groups))
})?;
let date: String = row.get(4)?;
let date = chrono::DateTime::parse_from_rfc3339(&date)
.expect("failed to get source date in Sqlite::list_sources");
let date = date.with_timezone(&Local);
Ok(admin::Source {
id,
name: row.get(1)?,
description: row.get(2)?,
url: row.get(3)?,
date,
files,
groups,
parent: row.get(5)?,
})
})? {
if let Ok(source) = result {
sources.push(source);
} else {
bail!("Failed to fetch source data");
}
}
Ok(sources)
}
#[cfg(any(test, feature = "admin"))]
pub fn create_source(
&self,
name: &str,
description: Option<&str>,
url: Option<&str>,
date: chrono::DateTime<Local>,
releasable: bool,
) -> Result<i32> {
let result = self.conn.execute(
"insert into source(name, description, url, firstacquisition, releasable) values (?1, ?2, ?3, ?4, ?5);",
params![&name, &description, url, date.to_rfc3339(), releasable],
)?;
if result != 1 {
bail!("failed to create source {name}");
}
let mut statement = self.conn.prepare("select id from source where name = ?1")?;
let gid = statement.query_row([name], |row| {
let sid: i32 = row.get(0)?;
Ok(sid)
})?;
Ok(gid)
}
#[cfg(any(test, feature = "admin"))]
pub fn edit_user(
&self,
uid: i32,
uname: &str,
fname: &str,
lname: &str,
email: &str,
) -> Result<()> {
self.conn.execute("update person set uname = ?1, email = ?2, firstname = ?3, lastname = ?4 where id = ?5;", params![&uname, &email, &fname, &lname, &uid])?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn deactivate_user(&self, uid: i32) -> Result<()> {
self.conn.execute(
"update person set password = null, apikey = null where id = ?1;",
params![&uid],
)?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn file_types_counts(&self) -> Result<HashMap<String, u32>> {
let mut types_counts = HashMap::default();
let mut statement = self.conn.prepare("SELECT filetype.name, count(file.id) from file join filetype on (file.filetypeid = filetype.id) group by 1")?;
let results = statement
.query_map([], |row| {
let name = row.get(0)?;
let count: i32 = row.get(1)?;
Ok((name, count))
})?
.flatten();
for (name, count) in results {
types_counts.insert(name, count as u32);
}
Ok(types_counts)
}
#[cfg(any(test, feature = "admin"))]
pub fn create_label(&self, name: &str, parent: Option<i64>) -> Result<u64> {
self.conn.execute(
"insert into label(name, parent) values (?1, ?2);",
params![&name, &parent],
)?;
let mut statement = self.conn.prepare("select id from label where name = ?1")?;
let lid = statement.query_row([name], |row| {
let lid: i64 = row.get(0)?;
Ok(lid)
})?;
Ok(lid as u64)
}
#[cfg(any(test, feature = "admin"))]
pub fn edit_label(&self, id: u64, name: &str, parent: Option<u64>) -> Result<()> {
self.conn.execute(
"update label set name = ?1, parent = ?2 where id = ?3",
params![&name, &parent, &id],
)?;
Ok(())
}
#[cfg(any(test, feature = "admin"))]
pub fn label_id_from_name(&self, name: &str) -> Result<u64> {
let mut statement = self.conn.prepare("select id from label where name = ?1")?;
let lid = statement.query_row([name], |row| {
let lid: i64 = row.get(0)?;
Ok(lid)
})?;
Ok(lid as u64)
}
}
unsafe impl Send for Sqlite {}
unsafe impl Sync for Sqlite {}
impl Display for Sqlite {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "SQLite client")
}
}
impl Debug for Sqlite {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "SQLite client")
}
}