From 38128975e652d8385fbc81578ec9ee462c5f096b Mon Sep 17 00:00:00 2001 From: joon Date: Sun, 24 Jul 2022 17:03:14 -0700 Subject: [PATCH] Updates and view --- wtfis/clients/passivetotal.py | 20 ++++------- wtfis/main.py | 30 +++++++++------- wtfis/models/passivetotal.py | 1 + wtfis/utils.py | 6 ++++ wtfis/view.py | 65 +++++++++++++++++++++++++++++++++++ 5 files changed, 97 insertions(+), 25 deletions(-) create mode 100644 wtfis/utils.py create mode 100644 wtfis/view.py diff --git a/wtfis/clients/passivetotal.py b/wtfis/clients/passivetotal.py index a8fbe08..222f471 100644 --- a/wtfis/clients/passivetotal.py +++ b/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)) diff --git a/wtfis/main.py b/wtfis/main.py index fc320bd..3d05e70 100644 --- a/wtfis/main.py +++ b/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"Reputation: {domain.reputation}")) - print(HTML(f"Registrar: {domain.registrar}")) - print(HTML(f"Last DNS Records Date: {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)) diff --git a/wtfis/models/passivetotal.py b/wtfis/models/passivetotal.py index 85c4505..a8695a5 100644 --- a/wtfis/models/passivetotal.py +++ b/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] diff --git a/wtfis/utils.py b/wtfis/utils.py new file mode 100644 index 0000000..e1715a9 --- /dev/null +++ b/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() diff --git a/wtfis/view.py b/wtfis/view.py new file mode 100644 index 0000000..88b9670 --- /dev/null +++ b/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)))