mysticmomba/AirlockTools
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Post Black Linting

Browse Source
This commit is contained in:
Zarithas
2025-11-06 11:04:59 -05:00
parent f33b041ac0
commit b538f12e9a
20 changed files with 1106 additions and 618 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/agent.py
+11 -15
View File
@@ -38,23 +38,19 @@ class Agent:
status_text: Optional[str] = field(default=None)
# Class-level status map
status_map: ClassVar[dict] = {
0: "Offline",
1: "Online",
2: "Hidden",
3: "Safemode"
}
status_map: ClassVar[dict] = {0: "Offline", 1: "Online", 2: "Hidden", 3: "Safemode"}
def enrich_with_policies(self, policies: List[Policy]):
"""Enrich the agent with groupname and human-readable status."""
self.status_text = self.status_map.get(self.status, "Unknown")
for policy in policies:
if policy.groupid == self.groupid:
self.groupname = policy.name
break
if not self.groupname:
self.groupname = "Unknown"
"""Enrich the agent with groupname and human-readable status."""
self.status_text = self.status_map.get(self.status, "Unknown")
for policy in policies:
if policy.groupid == self.groupid:
self.groupname = policy.name
break
if not self.groupname:
self.groupname = "Unknown"
"""
from models.agent import Agent
models/execution.py
+71 -36
View File
@@ -35,11 +35,13 @@ logger = logging.getLogger(__name__)
dotenv.load_dotenv()
@dataclass
class Hash:
"""
Hash model representing Hash data
"""
sha256: str
applications: str
baselines: str
@@ -61,7 +63,7 @@ class Hash:
sha384: str
sha512: str
at_decision: Optional[str] = None
def to_dict(self):
return asdict(self)
@@ -99,11 +101,15 @@ class Hash:
for hash_obj in hashes:
publisher = hash_obj.publisher or ""
description = hash_obj.description or ""
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
reputation = (
hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
)
scannermatch = reputation.get("scannermatch")
logger.debug(f"Evaluating hash: {hash_obj}")
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}")
logger.debug(
f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
)
# 1. Unapproved: bad publisher or PUP
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
@@ -127,9 +133,9 @@ class Hash:
# 3. Approved or Unapproved based on threat level
try:
score = int(scannermatch) # pyright: ignore[reportArgumentType]
score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}")
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
if score > threat_tolerance: # pyright: ignore[reportOperatorIssue]
logger.debug("Unapproved: Unsigned file with high threat score.")
hash_obj.at_decision = "unapproved"
unapproved_count += 1
@@ -138,15 +144,17 @@ class Hash:
hash_obj.at_decision = "approved"
approved_count += 1
except (ValueError, TypeError):
logger.debug("Needs Review: Scannermatch score is missing or invalid. — {e}")
logger.debug(
"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review"
needs_review_count += 1
logger.debug(f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}")
logger.debug(
f"Final counts — Needs Review: {needs_review_count}, Approved: {approved_count}, Unapproved: {unapproved_count}"
)
return hashes
@classmethod
def export_to_csv(cls, hash_list, directory_path):
"""
@@ -203,11 +211,12 @@ class ExecutionHistoryRecord:
localip: Optional[str] = None
extid: Optional[str] = None
extname: Optional[str] = None
exttype: Optional[int] = None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
exttype: Optional[int] = (
None # 1 = CRX Chromium Extension, 2 = XPI Firefox Extension
)
extbrowser: Optional[int] = None # 1 = Chrome, 2 = Firefox, 3 = Edge
hash_obj: Optional[Hash] = None
@classmethod
def from_dict(cls, data: dict):
mandatory_fields = [
@@ -224,7 +233,9 @@ class ExecutionHistoryRecord:
"datetime",
]
missing_fields = [
field for field in mandatory_fields if field not in data or data[field] is None
field
for field in mandatory_fields
if field not in data or data[field] is None
]
if missing_fields:
raise ValueError(f"Missing mandatory fields: {missing_fields}")
@@ -254,7 +265,7 @@ class ExecutionHistoryRecord:
extname=data.get("extname"),
exttype=data.get("exttype"),
extbrowser=data.get("extbrowser"),
hash_obj=data.get("hash_obj")
hash_obj=data.get("hash_obj"),
)
@classmethod
@@ -263,7 +274,9 @@ class ExecutionHistoryRecord:
) -> List["ExecutionHistoryRecord"]:
executions = []
for policy in selected_policies:
execs = airlock_libs.pull_policy_exec_histories(api, policy.name, str([1,2,6,7]), history_days)
execs = airlock_libs.pull_policy_exec_histories(
api, policy.name, str([1, 2, 6, 7]), history_days
)
if execs:
data = json.loads(execs)
exechistories = data.get("response", {}).get("exechistories", [])
@@ -274,8 +287,12 @@ class ExecutionHistoryRecord:
df = df.drop_duplicates(subset=["sha256", "filename", "hostname"])
df = df.sort_values(by=["sha256", "filename"])
executions.extend([cls.from_dict(row.to_dict()) for _, row in df.iterrows()])
logger.debug(f"Staging of Execution history for policy: {policy.name} is complete")
executions.extend(
[cls.from_dict(row.to_dict()) for _, row in df.iterrows()]
)
logger.debug(
f"Staging of Execution history for policy: {policy.name} is complete"
)
print(
colorText(
f"Staging of Execution history for policy: {policy.name} is complete",
@@ -284,20 +301,23 @@ class ExecutionHistoryRecord:
)
return executions
@staticmethod
def enrich_with_hashes(
api: AirlockAPIWrapper,
executions: List["ExecutionHistoryRecord"]
api: AirlockAPIWrapper, executions: List["ExecutionHistoryRecord"]
) -> List["ExecutionHistoryRecord"]:
"""
Enriches each ExecutionHistoryRecord with a matching Hash object by querying the API.
"""
sha256_list = list({e.sha256.strip().lower() for e in executions if e.sha256})
logger.info(f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records.")
logger.info(
f"Extracted {len(sha256_list)} unique sha256 values from {len(executions)} execution records."
)
if not sha256_list:
logger.warning("No sha256 values found in execution records. Skipping enrichment.")
logger.warning(
"No sha256 values found in execution records. Skipping enrichment."
)
return executions
logger.debug("Querying hash data from API...")
@@ -306,8 +326,10 @@ class ExecutionHistoryRecord:
hash_objects = []
required_fields = {
f.name for f in dataclasses.fields(Hash)
if f.default == dataclasses.MISSING and f.default_factory == dataclasses.MISSING
f.name
for f in dataclasses.fields(Hash)
if f.default == dataclasses.MISSING
and f.default_factory == dataclasses.MISSING
}
for sha256, (_, row) in zip(sha256_list, hash_df.iterrows()):
@@ -345,11 +367,15 @@ class ExecutionHistoryRecord:
exec_record.hash_obj = hash_obj
enriched_count += 1
logger.info(f"Enriched {enriched_count} out of {len(executions)} execution records with hash data.")
logger.info(
f"Enriched {enriched_count} out of {len(executions)} execution records with hash data."
)
return executions
@staticmethod
def categorize_executions_by_hash_decision(executions: List["ExecutionHistoryRecord"]) -> List["ExecutionHistoryRecord"]:
def categorize_executions_by_hash_decision(
executions: List["ExecutionHistoryRecord"],
) -> List["ExecutionHistoryRecord"]:
"""
Categorizes the hash_obj of each ExecutionHistoryRecord based on publisher, description, and reputation.
@@ -373,11 +399,15 @@ class ExecutionHistoryRecord:
publisher = hash_obj.publisher or ""
description = hash_obj.description or ""
reputation = hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
reputation = (
hash_obj.reputation if isinstance(hash_obj.reputation, dict) else {}
)
scannermatch = reputation.get("scannermatch")
logger.debug(f"Evaluating hash: {hash_obj}")
logger.debug(f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}")
logger.debug(
f"Publisher: {publisher}, Description: {description}, Scannermatch: {scannermatch}"
)
# 1. Unapproved: bad publisher or PUP
if re.search(bad_publishers_pattern, publisher, re.IGNORECASE):
@@ -401,7 +431,7 @@ class ExecutionHistoryRecord:
# 3. Approved or Unapproved based on threat level
try:
score = int(scannermatch) # pyright: ignore[reportArgumentType]
score = int(scannermatch) # pyright: ignore[reportArgumentType]
logger.debug(f"Parsed scannermatch score: {score}")
if threat_tolerance is not None and score >= threat_tolerance:
logger.debug("Unapproved: Unsigned file with high threat score.")
@@ -412,7 +442,9 @@ class ExecutionHistoryRecord:
hash_obj.at_decision = "approved"
approved_count += 1
except (ValueError, TypeError) as e:
logger.debug(f"Needs Review: Scannermatch score is missing or invalid. — {e}")
logger.debug(
f"Needs Review: Scannermatch score is missing or invalid. — {e}"
)
hash_obj.at_decision = "needs_review"
needs_review_count += 1
@@ -422,12 +454,14 @@ class ExecutionHistoryRecord:
)
return executions
@classmethod
def sort_by_hash_decision(
cls, executions: List["ExecutionHistoryRecord"]
) -> Tuple[List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"], List["ExecutionHistoryRecord"]]:
def sort_by_hash_decision(cls, executions: List["ExecutionHistoryRecord"]) -> Tuple[
List["ExecutionHistoryRecord"],
List["ExecutionHistoryRecord"],
List["ExecutionHistoryRecord"],
List["ExecutionHistoryRecord"],
]:
"""
Sorts ExecutionHistoryRecord objects into approved, unapproved, needs_review, and unknown groups
based on the value of hash_obj.at_decision.
@@ -453,7 +487,9 @@ class ExecutionHistoryRecord:
else:
unknown.append(record)
logger.info(f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:")
logger.info(
f"[ExecutionHistoryRecord] Sorted {len(sorted_executions)} records by hash_obj.at_decision:"
)
logger.info(f" Approved: {len(approved)}")
logger.info(f" Unapproved: {len(unapproved)}")
logger.info(f" Needs Review: {len(needs_review)}")
@@ -462,7 +498,6 @@ class ExecutionHistoryRecord:
return approved, unapproved, needs_review, unknown
"""
executions = ExecutionHistoryRecord.from_policies(api, selected_policies, type_=[0,1,3], history_days=30)
models/policy.py
+6 -2
View File
@@ -29,7 +29,9 @@ class Policy:
def __repr__(self):
# Show all current attributes, including dynamically added ones
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
attrs = ", ".join(
f"{key}={repr(value)}" for key, value in self.__dict__.items()
)
return f"<Execution({attrs})>"
def to_dict(self):
@@ -53,7 +55,9 @@ class Allowlist:
def __repr__(self):
# Show all current attributes, including dynamically added ones
attrs = ", ".join(f"{key}={repr(value)}" for key, value in self.__dict__.items())
attrs = ", ".join(
f"{key}={repr(value)}" for key, value in self.__dict__.items()
)
return f"<Execution({attrs})>"
def to_dict(self):