Merge pull request #15 from DominicTWHV/dev

Merge a few experimental/WIP code into main

Author: DominicTWHV

edge/database/driver.py

@@ -7,7 +7,7 @@
 
 from edge.logger.context import db_logger
 
-from edge.registry.database import PrimaryDBConf, SecurityDBConf
+from edge.registry.database import PrimaryDBConf, SecurityDBConf, MetadataDBConf
 
 class DBInitManager:
     @staticmethod
@@ -33,6 +33,7 @@ async def _init_sqlite(db_type: Literal["primary", "security_core", "security_ge
             "primary": (PrimaryDBConf.sqlite_schema_path),
             "security_core": (SecurityDBConf.sqlite_schema_path_core),
             "security_generic": (SecurityDBConf.sqlite_schema_path_generic),
+            "metadata": (MetadataDBConf.sqlite_schema_path),
         }
 
         schema_path = config_map[db_type]

edge/database/schema/metadata_tracker.sql

@@ -0,0 +1,14 @@
+CREATE TABLE IF NOT EXISTS `dataset_metadata` (
+    `repository` VARCHAR, -- github repository, ex: DominicTWHV/CTI-DB-DUMP | If using a txt dataset, the url will be stored if the dataset isn't on github.
+    `licensing` VARCHAR, -- dataset licensing
+    `dataset_type` VARCHAR, -- cockatoo_core, text (like traditional url blocklists)
+    `remote_update_interval` INT, -- in hours, how often the remote updates
+
+    `num_of_url_entries` INT,
+    `num_of_file_entries` INT,
+    `num_of_invite_entries` INT,
+
+    `enabled` INT DEFAULT 0, -- whether this dataset is enabled (1 = true, 0 = false)
+
+    `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP
+);

edge/database/schema/primary.sql

@@ -9,17 +9,6 @@ CREATE TABLE IF NOT EXISTS `configs` (
     `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP
 );
 
-CREATE TABLE IF NOT EXISTS `datasets` (
-    `security_repository` VARCHAR PRIMARY KEY, -- github repository for security dataset, ex: DominicTWHV/CTI-DB-DUMP
-
-    `in_use` INT DEFAULT 0, -- whether this dataset is actively being used (1 = true, 0 = false)
-    `locked` INT DEFAULT 0, -- whether this dataset is locked from deletion (manually)
-
-    `remote_update_interval` INT DEFAULT 0, -- how often in minutes is this dataset updates by the remote peer (set automatically, do not manually edit)
-    
-    `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP
-);
-
 CREATE TABLE IF NOT EXISTS `actions` (
     `server_id` VARCHAR PRIMARY KEY, -- server identifier
 

edge/database/schema/security_generic.sql

@@ -1,16 +1,3 @@
-CREATE TABLE IF NOT EXISTS `dataset_metadata` (
-    `repository` VARCHAR, -- github repository, ex: DominicTWHV/CTI-DB-DUMP | If using a txt dataset, the url will be stored if the dataset isn't on github.
-    `licensing` VARCHAR, -- dataset licensing
-    `dataset_type` VARCHAR, -- cockatoo_core, text (like traditional url blocklists)
-    `remote_update_interval` INT, -- in hours, how often the remote updates
-
-    `num_of_url_entries` INT,
-    `num_of_file_entries` INT,
-    `num_of_invite_entries` INT,
-
-    `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP
-);
-
 CREATE TABLE IF NOT EXISTS `url_cache` (
   `id` INTEGER PRIMARY KEY AUTOINCREMENT,
 
@@ -19,18 +6,24 @@ CREATE TABLE IF NOT EXISTS `url_cache` (
   `url_hash` VARCHAR(64) NOT NULL,
 
   `times_seen` INT DEFAULT 1, -- how many times this url has been seen by core
+
+  `mother_set_name` VARCHAR, -- the dataset this url originated from
 
   `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP -- auto timestamp for last seen time
 );
 
 CREATE TABLE IF NOT EXISTS `file_hash_cache` (
   `sha256` VARCHAR PRIMARY KEY,
 
+  `mother_set_name` VARCHAR, -- the dataset this url originated from
+
   `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP -- auto timestamp for last seen time
 );
 
 CREATE TABLE IF NOT EXISTS `invite_cache` (
   `invite` VARCHAR PRIMARY KEY, -- will be a sha256 hash
 
+  `mother_set_name` VARCHAR, -- the dataset this url originated from
+
   `last_update` timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP -- auto timestamp for last seen time
 );

edge/database/sqlite/metadata.db

@@ -4,7 +4,7 @@
 
 from edge.database.driver import SQLiteDriver, DBInitManager
 
-from edge.registry.database import PrimaryDBConf, SecurityDBConf
+from edge.registry.database import PrimaryDBConf, SecurityDBConf, MetadataDBConf
 
 class AutoDBMgr:
 
@@ -13,5 +13,9 @@ async def init_db() -> None: #initializes static db instances. Not used for data
 
         for db_path, db_type in [
             (PrimaryDBConf.sqlite_db_path, "primary"),
+            (SecurityDBConf.sqlite_db_path_core, "security_core"),
+            (SecurityDBConf.sqlite_db_path_generic, "security_generic"),
+            (MetadataDBConf.sqlite_db_path, "metadata"),
+            
         ]:
             await DBInitManager.init_db(db_type=db_type, db_path=db_path)

edge/database/sqlite/security_core.db

@@ -39,7 +39,7 @@ class DownloadManager:
 
     @staticmethod
     async def download_file(url: str, dos_protection: bool = SessionConfigs.dos_protection, max_size: int = SessionConfigs.max_file_size) -> str:
-        session = SessionFactory().grab_session()
+        session = await SessionFactory().grab_session()
         async with session.get(url) as response:
 
             if SessionConfigs.raise_for_status:

edge/database/sqlite/security_generic.db

@@ -0,0 +1,74 @@
+import magic # file type identification\
+import tldextract #domain extraction
+
+from edge.helper.downloadMgr import URLParser, DownloadManager
+from edge.helper.aiohttpSessionFactory import SessionFactory
+
+from edge.logger.context import networking_logger
+
+from edge.registry.networking import DatasetDownloadConfigs, FileSourceIdentConfigs
+
+class Ident:
+
+    @staticmethod
+    async def remote_file_type(url: str) -> str:
+        session = await SessionFactory().grab_session()
+        async with session.get(url) as response:
+            mime_header = response.headers.get('Content-Type', '').lower()
+
+            chunk = await response.content.read(DatasetDownloadConfigs.read_chunk_size)  #read a certain amount of bytes for identification
+
+            detected_type = magic.from_buffer(chunk, mime=True) #json, text/plain, etc.
+
+            networking_logger.debug(f"File Type Identification: Detected type: {detected_type}, Mime Header: {mime_header}")
+
+            return detected_type, mime_header
+        
+    @staticmethod
+    async def file_source(url: str) -> str:
+        extracted = tldextract.extract(url)
+        
+        domain = extracted.registered_domain.lower()
+
+        if domain in FileSourceIdentConfigs.github_domains:
+            networking_logger.debug(f"URL {url} identified as GitHub domain: {domain}")
+            return "github"
+        
+        networking_logger.debug(f"URL {url} source is non-GitHub domain: {domain}")
+        return "unknown"
+    
+class Verify:
+
+    @staticmethod
+    async def file_type_allowed(detected_type: str) -> bool:
+        if not DatasetDownloadConfigs.validate_file_type:
+            networking_logger.debug(f"File Type Verification: Validation disabled. Bypassing checks. File type: {detected_type}")
+            return True
+
+        if detected_type in DatasetDownloadConfigs.allowed_mime_types:
+            networking_logger.info(f"File Type Verification: Detected type {detected_type} is allowed.")
+            return True
+        
+        else:
+            networking_logger.warning(f"File Type Verification: Detected type {detected_type} is NOT allowed.")
+            return False
+    
+class DSDownload:
+
+    @staticmethod
+    async def pipeline(url: str) -> str:
+        detected_type, mime_header = await Ident.remote_file_type(url)
+        source = await Ident.file_source(url)
+
+        if not await Verify.file_type_allowed(detected_type):
+            networking_logger.error(f"Dataset Download Pipeline: File type {detected_type} not allowed. Aborting download for safety. URL: {url}")
+            return
+        
+        if source == "github": #parse out the raw github content url (raw.githubusercontent.com links will NOT be identified as github)
+            raw_url = await URLParser.parse_github_url(url)
+            url = raw_url.append("metadata.json") #append the standard dataset filename for Cockatoo Core dataset format
+
+            content = await DownloadManager.download_file(url) #download the mnetadata file into a variable
+
+        else:
+            content = await DownloadManager.download_file(url) # -> this is the set as not using cockatoo core dataset format