Updates and view

wtfis/clients/passivetotal.py

@@ -27,20 +27,14 @@ def _get(self, request: str, params: Optional[dict] = None) -> Optional[dict]:
         except (HTTPError, JSONDecodeError):
             raise
 
-    def passive(self, domain: str) -> dict:
+    def _query(self, path: str, query: str) -> Optional[dict]:
         return self._get(
-            "/dns/passive",
-            params={
-                "query": domain,
-            },
+            path,
+            params={"query": query}
         )
 
+    def get_passive_dns(self, domain: str) -> dict:
+        return self._query("/dns/passive", domain)
+
     def get_whois(self, domain: str) -> Optional[Whois]:
-        return Whois.parse_obj(
-            self._get(
-                "/whois",
-                params={
-                    "query": domain
-                }
-            )
-        )
+        return Whois.parse_obj(self._query("/whois", domain))

wtfis/main.py

@@ -1,17 +1,12 @@
-import datetime
 import os
 import sys
 
 from dotenv import load_dotenv
-from prompt_toolkit import HTML, print_formatted_text as print
 
 from wtfis.clients.passivetotal import PTClient
 from wtfis.clients.virustotal import VTClient
 from wtfis.models.virustotal import Domain
-
-
-def iso_date(unix_time: int) -> str:
-    return datetime.datetime.utcfromtimestamp(unix_time).isoformat()
+from wtfis.view import View
 
 
 def main():
@@ -22,11 +17,22 @@ def main():
     vt = VTClient(os.environ.get("VT_API_KEY"))
     domain = Domain.parse_obj(vt.get_domain(sys.argv[1]))
 
-    print(HTML(f"<b>Reputation:</b> {domain.reputation}"))
-    print(HTML(f"<b>Registrar:</b> {domain.registrar}"))
-    print(HTML(f"<b>Last DNS Records Date:</b> {iso_date(domain.last_dns_records_date)}"))
-
     pt = PTClient(os.environ.get("PT_API_USER"), os.environ.get("PT_API_KEY"))
-    passive = pt.get_whois(sys.argv[1])
+    whois = pt.get_whois(sys.argv[1])
+
+    console = View(whois, domain)
+    console.print()
+
+    # console = Console()
+
+    # top = Text(sys.argv[1], style="bold yellow", justify="center", end="\n")
+
+    # text = Text()
+    # text.append("Reputation: ", style="bold magenta")
+    # text.append(f"{domain.reputation}\n")
+    # text.append(f"Registrar: {domain.registrar}\n")
+    # text.append(f"Last DNS Records Date: {iso_date(domain.last_dns_records_date)}")
+
+    # text_group = Group(top, text)
 
-    print(passive)
+    # console.print(Panel(text_group, title="virustotal", expand=False))

wtfis/models/passivetotal.py

@@ -11,6 +11,7 @@ class Registrant(BaseModel):
 
 class Whois(BaseModel):
     contactEmail: str
+    domain: str
     expiresAt: str
     name: str
     nameServers: List[str]

wtfis/utils.py

@@ -0,0 +1,6 @@
+import datetime
+
+
+def iso_date(unix_time: int) -> str:
+    """ Convert epoch time to ISO 8601 """
+    return datetime.datetime.utcfromtimestamp(unix_time).isoformat()

wtfis/view.py

@@ -0,0 +1,65 @@
+from rich.columns import Columns
+from rich.console import Console, Group
+from rich.padding import Padding
+from rich.panel import Panel
+from rich.text import Text
+from typing import Optional
+
+from wtfis.models.passivetotal import Whois
+from wtfis.models.virustotal import Domain
+from wtfis.utils import iso_date
+
+
+class View:
+    """
+    Handles the look of the output
+    """
+    def __init__(self, whois: Whois, vt_domain: Domain):
+        self.console = Console()
+        self.whois = whois
+        self.vt = vt_domain
+
+    def _gen_heading_text(self, heading: str) -> Text:
+        heading_style = "bold yellow"
+        return Text(heading, style=heading_style, justify="center", end="\n")
+
+    def _gen_fv_text(self, *params) -> Text:
+        """ Each param should be a tuple of (field, value, value_style_overrides) """
+        field_style = "bold magenta"
+        text = Text()
+        for idx, item in enumerate(params):
+            value_style = None
+            if len(item) == 3:
+                field, value, value_style = item
+            else:
+                field, value = item
+            text.append(field + " ", style=field_style)
+            text.append(str(value), style=value_style)
+            if not idx == len(params) - 1:
+                text.append("\n")
+        return text
+
+    def _gen_panel(self, title: str, body: Text, heading: Optional[Text] = None) -> Panel:
+        renderable = Group(heading, body) if heading else body
+        return Panel(renderable, title=title, expand=False)
+
+    def whois_panel(self) -> Panel:
+        heading = self._gen_heading_text(self.whois.domain)
+        body = "foo"
+        return self._gen_panel("whois", body, heading)
+
+    def vt_panel(self) -> Panel:
+        heading = self._gen_heading_text("goo.bar.baz")
+        body = self._gen_fv_text(
+            ("Reputation:", self.vt.reputation),
+            ("Registrar:", self.vt.registrar),
+        )
+
+        return self._gen_panel("virustotal", body, heading)
+
+    def print(self):
+        renderables = [
+            self.vt_panel(),
+            self.whois_panel(),
+        ]
+        self.console.print(Padding(Columns(renderables), (1, 0)))