diff --git a/setup.cfg b/setup.cfg index f962f22..901db42 100644 --- a/setup.cfg +++ b/setup.cfg @@ -7,6 +7,7 @@ packages = find: install_requires = pydantic requests + rich tldextract [options.entry_points] diff --git a/wtfis/clients/virustotal.py b/wtfis/clients/virustotal.py index 08dc74f..e4f4eca 100644 --- a/wtfis/clients/virustotal.py +++ b/wtfis/clients/virustotal.py @@ -26,7 +26,7 @@ def _get(self, request: str) -> Optional[dict]: resp = self.s.get(self.baseurl + request) resp.raise_for_status() - return json.loads(json.dumps((resp.json())))["data"]["attributes"] + return json.loads(json.dumps((resp.json()))) except (HTTPError, JSONDecodeError): raise diff --git a/wtfis/models/passivetotal.py b/wtfis/models/passivetotal.py index a8695a5..4215d42 100644 --- a/wtfis/models/passivetotal.py +++ b/wtfis/models/passivetotal.py @@ -1,12 +1,12 @@ from pydantic import BaseModel -from typing import List +from typing import List, Optional class Registrant(BaseModel): - organization: str + organization: Optional[str] email: str name: str - telephone: str + telephone: Optional[str] class Whois(BaseModel): @@ -15,7 +15,7 @@ class Whois(BaseModel): expiresAt: str name: str nameServers: List[str] - organization: str + organization: Optional[str] registered: str registrant: Registrant registrar: str diff --git a/wtfis/models/virustotal.py b/wtfis/models/virustotal.py index 3f4fe7a..b0bb102 100644 --- a/wtfis/models/virustotal.py +++ b/wtfis/models/virustotal.py @@ -30,16 +30,13 @@ class PopularityRanks(BaseModel): __root__: Dict[str, Popularity] -class Domain(BaseModel): - """ - Essential VT domain fields - """ +class Attributes(BaseModel): creation_date: int - jarm: str + jarm: Optional[str] last_analysis_results: LastAnalysisResults last_analysis_stats: LastAnalysisStats last_dns_records_date: int - last_https_certificate_date: int + last_https_certificate_date: Optional[int] last_modification_date: int last_update_date: int popularity_ranks: PopularityRanks @@ -48,3 +45,19 @@ class Domain(BaseModel): tags: List[str] whois: str whois_date: Optional[int] + + +class Data(BaseModel): + attributes: Attributes + id_: str + type_: str + + class Config: + fields = { + "id_": "id", + "type_": "type", + } + + +class Domain(BaseModel): + data: Data diff --git a/wtfis/view.py b/wtfis/view.py index 88b9670..85d793c 100644 --- a/wtfis/view.py +++ b/wtfis/view.py @@ -2,12 +2,38 @@ from rich.console import Console, Group from rich.padding import Padding from rich.panel import Panel +from rich.table import Table from rich.text import Text -from typing import Optional +from typing import Any, Callable, Optional from wtfis.models.passivetotal import Whois -from wtfis.models.virustotal import Domain -from wtfis.utils import iso_date +from wtfis.models.virustotal import Domain, LastAnalysisStats, PopularityRanks +# from wtfis.utils import iso_date + + +class Theme: + panel_title = "bright_blue" + heading = "bold yellow" + table_field = "bold bright_magenta" + table_value = "none" + inline_stat = "cyan" + info = "bold green" + warn = "bold yellow" + error = "bold red" + + @classmethod + def _get_theme_vars(cls): + return [ + attr for attr in dir(cls) + if not callable(getattr(cls, attr)) + and not attr.startswith("__") + ] + + def __init__(self, nocolor: Optional[bool] = False): + for attr in type(self)._get_theme_vars(): + value = getattr(self, attr) if not nocolor else "none" + setattr(self, attr, value) + self.nocolor = nocolor class View: @@ -18,48 +44,108 @@ def __init__(self, whois: Whois, vt_domain: Domain): self.console = Console() self.whois = whois self.vt = vt_domain + self.theme = Theme() def _gen_heading_text(self, heading: str) -> Text: - heading_style = "bold yellow" - return Text(heading, style=heading_style, justify="center", end="\n") + return Text(heading, style=self.theme.heading, 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_table(self, *params) -> Table: + """ Each param should be a tuple of (field, value) """ + # Set up table + grid = Table.grid(expand=False, padding=(0, 1)) + grid.add_column(style=self.theme.table_field) # Field + grid.add_column(style=self.theme.table_value, max_width=60) # Value + + # Populate rows + for item in params: + field, value = item + if value is None: # Skip if no value + continue + grid.add_row(field, value) + + return grid def _gen_panel(self, title: str, body: Text, heading: Optional[Text] = None) -> Panel: + panel_title = Text(title, style=self.theme.panel_title) renderable = Group(heading, body) if heading else body - return Panel(renderable, title=title, expand=False) + return Panel(renderable, title=panel_title, expand=False) + + def _cond_style(self, item: Any, func: Callable) -> Optional[str]: + """ Conditional style """ + return func(item) if not self.theme.nocolor else "none" + + def _gen_vt_analysis_stats(self, stats: LastAnalysisStats) -> Text: + # Custom style + def stats_style(malicious): + return self.theme.error if malicious >= 1 else self.theme.info + + total = stats.harmless + stats.malicious + stats.suspicious + stats.timeout + stats.undetected + return Text(f"{stats.malicious}/{total} malicious", style=self._cond_style(stats.malicious, stats_style)) + + def _gen_vt_reputation(self, reputation: int) -> Text: + # Custom style + def rep_style(reputation): + if reputation > 0: + return self.theme.info + elif reputation < 0: + return self.theme.error + + return Text(str(reputation), style=self._cond_style(reputation, rep_style)) + + def _gen_vt_popularity(self, popularity_ranks: PopularityRanks) -> Optional[Text]: + if len(popularity_ranks.__root__) == 0: + return None + + text = Text() + for source, popularity in popularity_ranks.__root__.items(): + text.append(f"{source} (") + text.append(str(popularity.rank), style=self.theme.inline_stat) + text.append(")") + if source != list(popularity_ranks.__root__.keys())[-1]: + text.append("\n") + return text def whois_panel(self) -> Panel: heading = self._gen_heading_text(self.whois.domain) - body = "foo" + body = self._gen_table( + ("Registrar:", self.whois.registrar), + ("Organization:", self.whois.organization), + ("Name:", self.whois.name), + ("E-mail:", self.whois.contactEmail), + ("Phone:", self.whois.registrant.telephone), + ) 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), - ) + attributes = self.vt.data.attributes + # Analysis + analysis = self._gen_vt_analysis_stats(attributes.last_analysis_stats) + + # Reputation + reputation = self._gen_vt_reputation(attributes.reputation) + + # Popularity + popularity = self._gen_vt_popularity(attributes.popularity_ranks) + + # Reputation style + def rep_style(reputation): + if reputation > 0: + return self.theme.info + elif attributes.reputation < 0: + return self.theme.error + + heading = self._gen_heading_text(self.vt.data.id_) + body = self._gen_table( + ("Analysis:", analysis), + ("Reputation:", reputation), + ("Popularity:", popularity), + ) return self._gen_panel("virustotal", body, heading) def print(self): renderables = [ - self.vt_panel(), self.whois_panel(), + self.vt_panel(), ] self.console.print(Padding(Columns(renderables), (1, 0)))