Minimum Working Product 0.6 💝

Mrida/Mrida/Mrida.vcxproj.filters

@@ -144,7 +144,7 @@
       <Filter>Header Files\trend</Filter>
     </ClInclude>
     <ClInclude Include="threat_database.h">
-      <Filter>Header Files</Filter>
+      <Filter>Header Files\database</Filter>
     </ClInclude>
   </ItemGroup>
 </Project>

Mrida/Mrida/mrida.cpp

@@ -226,7 +226,7 @@ int main(int argc, char** argv)
 			long min_size = std::stoll(req.get_param_value("min_size"));
 			long max_size = std::stoll(req.get_param_value("max_size"));
 			std::string type = req.get_param_value("type");
-			trendcpp tlsh;
+			threat_database tlsh;
 			long id = tlsh.matching_hash_from_threat_db(tlsh_hash, type, min_size, max_size);
 			res.set_content(return_json(id), "application/json");
 		}
@@ -239,6 +239,21 @@ int main(int argc, char** argv)
 		res.set_content(send_success_response(), "application/json");
 	});
 
+	// Add threat
+	server.Post("/add_threat", [](const httplib::Request& req, httplib::Response& res) {
+		std::string threat_hash = req.get_param_value("hash");
+		std::string name = req.get_param_value("name");
+		unsigned long size = std::stol(req.get_param_value("size"));
+		std::string type = req.get_param_value("type");
+		threat_database database;
+		database.add_threat_to_database(threat_hash, name, size, type);
+		print_terminal_info();
+		set_terminal_color(CYAN);
+		std::cout << "ADDED HASH\n";
+		set_terminal_color();
+		res.set_content(send_success_response(), "application/json");
+	});
+
 	print_terminal_info();
 	std::cout << "Server started on: " << "127.0.0.1:" << 5660 << "\n";
 	server.listen("127.0.0.1", 5660);

Mrida/Mrida/threat_database.cpp

@@ -10,6 +10,7 @@
 
 threat_database::threat_database()
 {
+
 }
 
 
@@ -92,3 +93,63 @@ void threat_database::refactor()
 	// Reset the terminal color to original
 	set_terminal_color();
 }
+
+void threat_database::add_threat_to_database(std::string tlsh_hash, std::string threat_name, unsigned long file_size, std::string file_type)
+{
+	unsigned int file_type_id = mime_to_id(file_type);
+	try {
+		sqlite::database _threat_database("threat_db.db");
+		_threat_database << "create table if not exists threat(id unsigned bigint primary key, threat_hash text, threat_name text, threat_size unsigned int, threat_type unsigned int);";
+		unsigned long max_id = 0;
+		_threat_database << "select max(id) from threat" >> max_id;
+		max_id++;
+		_threat_database << "insert into threat(id, threat_hash, threat_name, threat_size, threat_type) values(?, ?, ?, ?, ?)" << max_id << tlsh_hash << threat_name << file_size << file_type_id;
+	}
+	catch (std::exception &e)
+	{
+		std::cout << e.what();
+	}
+}
+
+
+unsigned int threat_database::mime_to_id(std::string mime_type)
+{
+
+	sqlite::database db("threat_db.db");
+	db << "create table if not exists mime_table(mime text, id int)";
+	int count = 0;
+	db << "select count(id) from mime_table where mime=?" << mime_type >> count;
+	int max = 0;
+	db << "select max(id) from mime_table limit 1" >> max;
+	if (count == 0)
+	{
+		max++;
+		db << "insert into mime_table(mime, id) values(?, ?)" << mime_type << max;
+		return max;
+	}
+	else
+	{
+		unsigned int id;
+		db << "select id from mime_table where mime=? limit 1" << mime_type >> id;
+		return id;
+	}
+	return 0;
+}
+
+long threat_database::matching_hash_from_threat_db(std::string tlsh_hash, std::string file_type, long file_size_minimum, unsigned long file_size_maximum)
+{
+	long matched_id = -1;
+	sqlite::database threat_table("threat_db.db");
+	threat_table << "create table if not exists threat(id unsigned bigint primary key, threat_hash text, threat_name text, threat_size unsigned int, threat_type unsigned int);";
+	unsigned int file_id = mime_to_id(file_type);
+	threat_table << "select id, threat_hash from threat where threat_size>=? and threat_size<=? and threat_type=?"
+		<< file_size_minimum << file_size_maximum << file_id >> [&](unsigned long id, std::string threat_hash)
+	{
+		trendcpp trend;
+		if (trend.similarity_distance(tlsh_hash, threat_hash) < 20)
+		{
+			matched_id = id;
+		}
+	};
+	return matched_id;
+}

Mrida/Mrida/threat_database.h

@@ -1,6 +1,7 @@
 // SWAMI KARUPPASWAMI THUNNAI
 
 #pragma once
+#include "httplib.h"
 #include <iostream>
 #include <sqlite_modern_cpp.h>
 
@@ -12,5 +13,14 @@ class threat_database
 
 	// Refactor the threat database -- will remove duplicates in the threat database
 	void refactor();
+
+	// Bool add threat to database
+	void add_threat_to_database(std::string tlsh_hash, std::string threat_name, unsigned long file_size, std::string file_type);
+
+	// Mime Type to Id
+	unsigned int mime_to_id(std::string mime_type);
+
+	// Will get the similar hash matching id from threat db [returns -1 if nothing is matching]
+	long matching_hash_from_threat_db(std::string tlsh_hash, std::string file_type, long file_size_minimum, unsigned long file_size_maximum);
 };
 

Mrida/Mrida/trendcpp.cpp

@@ -160,19 +160,6 @@ const Tlsh * trendcpp::hash_file(std::string file_location)
 	return &th;
 }
 
-void trendcpp::add_threat_to_database(unsigned long int id, std::string tlsh_hash, std::string threat_name, unsigned long file_size, unsigned int file_type)
-{
-	try {
-		sqlite::database threat_database("threat_db.db");
-		threat_database << "create table if not exists threat(id unsigned bigint primary key, threat_hash text, threat_name text, threat_size unsigned int, threat_type unsigned int);";
-		threat_database << "insert into threat(id, threat_hash, threat_name, threat_size, threat_type) values(?, ?, ?, ?, ?)" << id << tlsh_hash << threat_name << file_size << file_type;
-	}
-	catch (std::exception &e)
-	{
-		std::cout << e.what();
-	}
-}
-
 int trendcpp::similarity_distance(std::string hash_one, std::string hash_two)
 {
 	Tlsh t1;
@@ -185,43 +172,3 @@ int trendcpp::similarity_distance(std::string hash_one, std::string hash_two)
 	return t1.totalDiff(&t2);
 }
 
-unsigned int trendcpp::mime_to_id(std::string mime_type)
-{
-
-	sqlite::database db("threat_db.db");
-	db << "create table if not exists mime_table(mime text, id int)";
-	int count = 0;
-	db << "select count(id) from mime_table where mime=?" << mime_type >> count;
-	int max = 0;
-	db << "select max(id) from mime_table limit 1" >> max;
-	if (count == 0)
-	{
-		max++;
-		db << "insert into mime_table(mime, id) values(?, ?)" << mime_type << max;
-		return max;
-	}
-	else
-	{
-		unsigned int id;
-		db << "select id from mime_table where mime=? limit 1" << mime_type >> id;
-		return id;
-	}
-	return 0;
-}
-
-long trendcpp::matching_hash_from_threat_db(std::string tlsh_hash, std::string file_type, long file_size_minimum, unsigned long file_size_maximum)
-{
-	long matched_id = -1;
-	sqlite::database threat_table("threat_db.db");
-	threat_table << "create table if not exists threat(id unsigned bigint primary key, threat_hash text, threat_name text, threat_size unsigned int, threat_type unsigned int);";
-	unsigned int file_id = mime_to_id(file_type);
-	threat_table << "select id, threat_hash from threat where threat_size>=? and threat_size<=? and threat_type=?"
-		<< file_size_minimum << file_size_maximum << file_id >> [&] (unsigned long id, std::string threat_hash)
-		{
-			if (similarity_distance(tlsh_hash, threat_hash) < 20)
-			{
-				matched_id = id;
-			}
-		};
-	return matched_id;
-}

Mrida/Mrida/trendcpp.h

@@ -16,16 +16,7 @@ class trendcpp
 	// Get TLSH object
 	const Tlsh* hash_file(std::string file_location);
 
-	// Bool add threat to database
-	void add_threat_to_database(unsigned long int id, std::string tlsh_hash, std::string threat_name, unsigned long file_size, unsigned int file_type);
-
 	// Getting the similarity distance
 	int similarity_distance(std::string hash_one, std::string hash_two);
-
-	// Mime Type to Id
-	unsigned int mime_to_id(std::string mime_type);
-
-	// Will get the similar hash matching id from threat db [returns -1 if nothing is matching]
-	long matching_hash_from_threat_db(std::string tlsh_hash, std::string file_type, long file_size_minimum, unsigned long file_size_maximum);
 };
 

mrida_gui/mrida.py

@@ -4,13 +4,15 @@
 from PyQt5.QtWidgets import QMainWindow, QApplication, QTabWidget
 from PyQt5.QtCore import Qt
 from scan import ScanWidget
+from update import UpdateWidget
 
 
 class MridaMainWidget(QTabWidget):
 
     def __init__(self):
         super().__init__()
         self.addTab(ScanWidget(), "Scan")
+        self.addTab(UpdateWidget(), "Update")
 
 
 class MridaMainWindow(QMainWindow):