misc: refactor code structure

lib/structure.py

@@ -87,14 +87,14 @@ class AbuseObject(BaseModel):
 
 
 class VTLastAnalysisStats(BaseModel):
-    harmless: int = 0
+    harmless: Optional[int] = 0
     type_unsupported: Optional[int] = None  # Field(..., alias="type-unsupported")
-    suspicious: int = 0
+    suspicious: Optional[int] = 0
     confirmed_timeout: Optional[int] = None  # Field(..., alias="confirmed-timeout")
-    timeout: int = 0
+    timeout: Optional[int] = 0
     failure: Optional[int] = None
-    malicious: int = 0
-    undetected: int = 0
+    malicious: Optional[int] = 0
+    undetected: Optional[int] = 0
 
 
 class VTSignatureInfo(BaseModel):

lib/worker.py

@@ -3,7 +3,13 @@
 import threading
 import pycurl
 from .util import resource_path
-from .structure import AbuseObject, VirusTotalObject, ShodanObject, NISTObject, CirclCVEObject
+from .structure import (
+    AbuseObject,
+    VirusTotalObject,
+    ShodanObject,
+    NISTObject,
+    CirclCVEObject,
+)
 import ouilookup
 import ipaddress
 import csv
@@ -208,7 +214,7 @@ def query(self, id, s):
         import base64
 
         try:
-            result = base64.b64decode(s.encode("utf-8"))
+            result = base64.b64decode(s, validate=True)
             decodedText = result.decode("utf-8")  # will add detection later
         except Exception as e:
             print(f"[base64decoder] encounter error: {e}")
@@ -251,7 +257,7 @@ def query(self, id, text, maxAge=90):
             "Key": self.apiKey,
             "Accept": "application/json",
         }
-        url = f"https://api.abuseipdb.com/api/v2/check?ipAddress={text}&maxAgeInDays={maxAge}"
+        url = f"https://api.abuseipdb.com/api/v2/check?ipAddress={text}&verbose&maxAgeInDays={maxAge}"
         try:
             if self.internetConfig[0] is not None:
                 self.timeout(id, text, 4000)
@@ -487,6 +493,7 @@ def query(self, id, cve, options={}):
             print(f"[worker-circlcve] error: {e}")
             self.ui.render(source="cve", box=(id, cve, None))
 
+
 class MacAddress:
     def __init__(self, ui):
         self.ui = ui
@@ -600,7 +607,7 @@ def run(self, id, target={}, text="", img=None):
             elif t == "cve":
                 # disable nist cve since they are slow and require api token now
                 # self.nistCVE.query(id, text)
-                self. circleCVE.query(id, text)
+                self.circleCVE.query(id, text)
             elif t == "netuser":
                 self.netUser.query(id, text)
             elif t == "base64":

main.py

@@ -9,13 +9,14 @@
 widget.set_appearance_mode("dark")
 
 
+app = DTSToolBox()
+
+
 def sigint_handler(sig, frame):
     app.exit_gracefully()
     sys.exit(-1)
 
 
 signal.signal(signal.SIGINT, sigint_handler)
 
-
-app = DTSToolBox()
 app.run()

widgets/custom.py → widgets/common.py

@@ -16,6 +16,7 @@ def __init__(
         web_btn=False,
         copy_btn=True,
         analyze_btn=False,
+        direct_btn=False,
         max_width=400,
         **kwargs,
     ):
@@ -39,7 +40,7 @@ def __init__(
         self.abtn = None
 
         if copy_btn:
-            icpy = Image.open(resource_path("lib/copy.png"))
+            icpy = Image.open(resource_path("lib/icons/copy.png"))
             self.cbtn = widget.CTkButton(
                 self,
                 text="",
@@ -48,7 +49,7 @@ def __init__(
                 image=widget.CTkImage(dark_image=icpy, light_image=icpy, size=(15, 15)),
             )
         if web_btn:
-            iweb = Image.open(resource_path("lib/web.png"))
+            iweb = Image.open(resource_path("lib/icons/web.png"))
             self.wbtn = widget.CTkButton(
                 self,
                 text="",
@@ -57,7 +58,7 @@ def __init__(
                 image=widget.CTkImage(dark_image=iweb, light_image=iweb, size=(15, 15)),
             )
         if analyze_btn:
-            ianalyze = Image.open(resource_path("lib/analyze.png"))
+            ianalyze = Image.open(resource_path("lib/icons/analyze.png"))
             self.abtn = widget.CTkButton(
                 self,
                 text="",
@@ -121,6 +122,14 @@ def clear(self):
             self.abtn.grid_remove()
 
 
+class DTSButton(widget.CTkButton):
+    def __init__(self, master, **kwargs):
+        super().__init__(master, self.cb_on_click, **kwargs)
+
+    def cb_on_click(self):
+        pass
+
+
 class DTSHistory(CTkListbox):
     def __init__(self, master, mainUI, **kwargs):
         super().__init__(master, command=self.cb_on_click, **kwargs)

widgets/report.py

@@ -1,6 +1,9 @@
 import customtkinter as widget
-from lib.tkdial import Meter
+from PIL import Image
 from iso3166 import countries
+from collections import Counter
+
+from lib.tkdial import Meter
 from lib.util import resource_path, unique
 from lib.structure import (
     AbuseObject,
@@ -10,9 +13,7 @@
     CirclCVEObject,
     DTSInputSource,
 )
-
-from .custom import DTSLabelWithBtn
-from PIL import Image
+from widgets.common import DTSLabelWithBtn
 
 
 class DTSGenericReport(widget.CTkFrame):
@@ -275,6 +276,14 @@ def populate(self, data: AbuseObject | None):
         if data.data.countryCode != "null":
             country = countries.get(data.data.countryCode)
             self.country.set("Country", f"{country.name}")
+
+        if data.data.abuseConfidenceScore != 0 and data.data.reports is not None:
+            categories = []
+            for r in data.data.reports:
+                categories += r.categories
+
+            c = Counter(categories)
+            print(c)
         self.error = False
 
 
@@ -552,10 +561,13 @@ def cb_on_analyze(self, event):
             source=DTSInputSource.TEXT_REPORT, text=self.textContent.get("0.0", "end")
         )
 
-    def populate(self, result: str, title="Text Report"):
+    def populate(self, data, title="Text Report"):
         self.title.configure(text=title)
         self.clear()
-        self.textContent.insert("0.0", result)
+        if data is not None:
+            self.textContent.insert("0.0", data)
+        else:
+            self.textContent.insert("0.0", "[An error happened]")
 
     def clear(self):
         self.textContent.delete("0.0", "end")

widgets/tabview.py

@@ -9,7 +9,7 @@
     DTSNISTCVEReport,
 )
 
-from widgets.custom import DTSLoading, DTSHistory, DTSLog
+from widgets.common import DTSLoading, DTSHistory, DTSLog
 from widgets.preferences import DTSPreferences
 
 from lib.analyzer import DTSAnalyzer
@@ -183,7 +183,6 @@ def render_from_worker(self, source, originalText, data):
             title = f"Report for {source.upper()}"
 
             self.textBoxData.delete("0.0", "end")
-            self.textBoxData.insert("0.0", "Nothing to show here ¯\_(ツ)_/¯")
 
             self.hide_other_reports(except_for="text")
             self.reports["text"].populate(data, title=title)