Skip to content

Commit

Permalink
Code refactoring for PEP8 compliance
Browse files Browse the repository at this point in the history
  • Loading branch information
@vincenzocaputo
vincenzocaputo committed Mar 31, 2024
1 parent 4132e38 commit 0f1e1e1
Showing 1 changed file with 39 additions and 58 deletions.
97 changes: 39 additions & 58 deletions pystixview/pystixview.py
View file
Original file line number Diff line number Diff line change
@@ -1,13 +1,18 @@
from pyvis.network import Network

from IPython.display import HTML
import os
import json
import base64
import warnings
import re

from pathlib import Path
from bs4 import BeautifulSoup

from IPython.display import HTML
from pyvis.network import Network

from stix2 import parsing
from pathlib import Path
from stix2.v21 import Bundle
from stix2.v21 import (
Bundle,
TLP_AMBER,
TLP_GREEN,
TLP_RED,
Expand Down Expand Up @@ -45,11 +50,6 @@
ThreatActor,
Tool,
Vulnerability)
import os
import json
import base64
import warnings
import re


class PySTIXView:
Expand Down Expand Up @@ -86,13 +86,18 @@ class PySTIXView:
}
}

__STYLES = ['square-flat',
'square-dark',
'square-lite',
'noback-dark',
'noback-flat',
'round-flat']

def __init__(self, height: str, width: str, notebook: bool = False,
select_menu: bool = False, filter_menu: bool = False,
style: str = 'square-flat'):

self.__notebook = notebook
self.__height = height
self.__width = width
if notebook:
self.__network = Network(height, width, directed=True,
notebook=notebook,
Expand All @@ -106,18 +111,12 @@ def __init__(self, height: str, width: str, notebook: bool = False,
filter_menu=filter_menu)
self.__icons_path = Path(os.path.dirname(__file__)) / 'icons'

STYLES = ['square-flat',
'square-dark',
'square-lite',
'noback-dark',
'noback-flat',
'round-flat']

if style in STYLES:
if style in self.__STYLES:
self.__style = style
else:
raise ValueError(f"Invalid style. \
Select from the following: {', '.join(STYLES)}")
Select from the following: {', '.join(self.__STYLES)}")

self.__network.barnes_hut(gravity=-5000,
central_gravity=0,
Expand Down Expand Up @@ -182,10 +181,9 @@ def __get_stix_object_type(self, object_to_test) -> str:
return "observable"

if isinstance(object_to_test, MarkingDefinition):
if object_to_test['id'] in self.__TLP_MARKINGS.keys():
if object_to_test['id'] in self.__TLP_MARKINGS:
return self.__TLP_MARKINGS[object_to_test['id']]['label']
else:
return "marking-definition"
return "marking-definition"

return None

Expand Down Expand Up @@ -236,7 +234,7 @@ def add_custom_stix_type(self, custom_type: str,
is already defined.
"""

if custom_type not in self.__custom_types.keys():
if custom_type not in self.__custom_types:
if node_icon:
if isinstance(node_icon, str):
if node_icon.startswith('http'):
Expand All @@ -253,7 +251,7 @@ def add_custom_stix_type(self, custom_type: str,
raise TypeError("Provide a valid color in hex rgb format")
self.__custom_types[custom_type]['label_name'] = label_name
else:
raise Exception(f"The custom type {custom_type}\
raise ValueError(f"The custom type {custom_type}\
is already defined")

def add_node(self,
Expand All @@ -264,27 +262,19 @@ def add_node(self,
Tool | Vulnerability | MarkingDefinition | AutonomousSystem |
DomainName | EmailAddress | EmailMessage | File |
IPv4Address | IPv6Address | MACAddress | NetworkTraffic |
URL | UserAccount | str | dict,
is_custom: bool = True,
node_icon: str = None,
color: str = None) -> bool:
URL | UserAccount | str | dict) -> bool:
"""Add a node to the graph
:param stix_obj: STIX Object (SDO, Observable or MarkingDefinition
to add to the graph
:param is_custom: Set to True to add a custom STIX Object
(one of node_icon or color must be provided)
:param node_icon: URLs or local path to the image to use as node icon
:param color: Color to be assigned to the node in place of the image.
Hex RGB format is accepted
:return: True if the node was added correctly
:raises KeyError: If a custom type does not have any
image or color for the node
:raises TypeError: If an invalid STIX Domain Object is provided
"""

node_img = None
if isinstance(stix_obj, dict) or isinstance(stix_obj, str):
if isinstance(stix_obj, (dict, str)):
stix_obj = parsing.parse(stix_obj, allow_custom=True)
else:
if not hasattr(stix_obj, 'type'):
Expand All @@ -294,7 +284,7 @@ def add_node(self,
label_name = 'name'
if not stix_object_type:
stix_type = stix_obj['type']
if stix_type in self.__custom_types.keys():
if stix_type in self.__custom_types:
if 'image' in self.__custom_types[stix_type].keys():
node_shape = "image"
node_img = self.__custom_types[stix_type]['image']
Expand Down Expand Up @@ -353,14 +343,13 @@ def add_node(self,
stix=node_title,
**stix_obj)
return True
else:
self.__network.add_node(node_id,
label=node_label,
shape=node_shape,
color=node_color,
stix=node_title,
**stix_obj)
return True
self.__network.add_node(node_id,
label=node_label,
shape=node_shape,
color=node_color,
stix=node_title,
**stix_obj)
return True

def add_bundle(self, bundle: Bundle | dict | str) -> bool:
"""Add a Bundle to the graph
Expand Down Expand Up @@ -421,8 +410,7 @@ def add_relationship(self, relationship: Relationship |

def _generate_graph(self, show_physics_buttons: bool = False,
show_node_buttons: bool = False,
show_edge_buttons: bool = False,
graph_option: str = None) -> str:
show_edge_buttons: bool = False) -> str:
"""Generate and return HTML code to render the graph.
In case of Jupyter Notebook, the graph is rendered
via IPython.display.HTML.
Expand All @@ -433,7 +421,6 @@ def _generate_graph(self, show_physics_buttons: bool = False,
options menu
:param show_edge_buttons: Set to True to show graph edge
options menu
:param graph_option:
:return: HTML code representin the graph.
If execute in a Jupyter Notebook,
an IPython.display.HTML object is returned
Expand All @@ -454,7 +441,7 @@ def _generate_graph(self, show_physics_buttons: bool = False,
html_graph = self.__network.generate_html(name)
bhtml = BeautifulSoup(html_graph, 'html.parser')
div_tag = bhtml.new_tag("div")
div_tag['id'] = "code_section"
div_tag['id'] = "code_section"
div_tag['style'] = "position: absolute;"
div_tag['style'] += "width: 30%;"
div_tag['style'] += "font-family: monospace;"
Expand All @@ -467,7 +454,6 @@ def _generate_graph(self, show_physics_buttons: bool = False,
div_tag['style'] += "border: 1px solid black;"
div_tag['style'] += "resize: both;"
div_tag['style'] += "overflow: auto;"

pre_tag = bhtml.new_tag("pre")
pre_tag['style'] = "width: 100%; height: 100%;"
pre_tag['style'] += "white-space: pre;"
Expand Down Expand Up @@ -497,8 +483,7 @@ def _generate_graph(self, show_physics_buttons: bool = False,

def show_graph(self, show_physics_buttons: bool = False,
show_node_buttons: bool = False,
show_edge_buttons: bool = False,
graph_option: str = None) -> str:
show_edge_buttons: bool = False) -> str:
"""Generate and return HTML code to render the graph.
In case of Jupyter Notebook, the graph is rendered
via IPython.display.HTML.
Expand All @@ -509,7 +494,6 @@ def show_graph(self, show_physics_buttons: bool = False,
options menu
:param show_edge_buttons: Set to True to show graph edge
options menu
:param graph_option:
:return: HTML code representin the graph.
If execute in a Jupyter Notebook,
an IPython.display.HTML object is returned
Expand All @@ -520,13 +504,11 @@ def show_graph(self, show_physics_buttons: bool = False,
show_edge_buttons)
if self.__notebook:
return HTML(html_graph)
else:
return html_graph
return html_graph

def save_graph(self, name, show_physics_buttons: bool = False,
show_node_buttons: bool = False,
show_edge_buttons: bool = False,
graph_option: str = None):
show_edge_buttons: bool = False):
"""Generate and save HTML file containing the graph.
:param name: Name of the file to save the graph as
Expand All @@ -536,13 +518,12 @@ def save_graph(self, name, show_physics_buttons: bool = False,
node options menu
:param show_edge_buttons: Set to True to show graph
edge options menu
:param graph_option:
"""

html_code = self._generate_graph(show_physics_buttons,
show_node_buttons,
show_edge_buttons)
with open(name, 'w') as fd:
with open(name, 'w', encoding="utf-8") as fd:
fd.write(html_code)

def to_json(self) -> str:
Expand Down

0 comments on commit 0f1e1e1

Please sign in to comment.