diff --git a/.pylintrc b/.pylintrc index b2465fb7fc..aa28c53036 100644 --- a/.pylintrc +++ b/.pylintrc @@ -673,7 +673,7 @@ dummy-variables-rgx=_+$|(_[a-zA-Z0-9_]*[a-zA-Z0-9]+?$)|dummy|^ignored_|^unused_ # Argument names that match this expression will be ignored. Default to name # with leading underscore. -ignored-argument-names=_.*|^ignored_|^unused_ +ignored-argument-names=_.*|^ignored_|^unused_|^ctx$ # Tells whether we should check for unused import in __init__ files. init-import=no diff --git a/api_client/python/timesketch_api_client/aggregation.py b/api_client/python/timesketch_api_client/aggregation.py index d09c8cc469..f44429f50d 100644 --- a/api_client/python/timesketch_api_client/aggregation.py +++ b/api_client/python/timesketch_api_client/aggregation.py @@ -13,11 +13,19 @@ # limitations under the License. """Timesketch API client library.""" +from __future__ import annotations + import datetime import getpass import json import logging + from typing import Any +from typing import Dict +from typing import Generator +from typing import List +from typing import Optional +from typing import TYPE_CHECKING import altair import pandas @@ -25,6 +33,9 @@ from . import error from . import resource +if TYPE_CHECKING: + from . import sketch as sketch_lib + logger = logging.getLogger("timesketch_api.aggregation") @@ -42,9 +53,14 @@ class Aggregation(resource.SketchResource): saved search. """ - resource_data: dict[str, Any] + resource_data: Dict[str, Any] + + def __init__(self, sketch: sketch_lib.Sketch) -> None: + """Initializes the Aggregation object. - def __init__(self, sketch): + Args: + sketch: An instance of Sketch object. + """ self._created_at = "" self._name = "" self._parameters = {} @@ -60,11 +76,13 @@ def __init__(self, sketch): super().__init__(sketch=sketch, resource_uri=resource_uri) @property - def created_at(self): + def created_at(self) -> str: """Returns a timestamp when the aggregation was created.""" return self._created_at - def _get_aggregation_buckets(self, entry, name=""): + def _get_aggregation_buckets( + self, entry: Dict[str, Any], name: str = "" + ) -> Generator[Dict[str, Any], None, None]: """Yields all buckets from an aggregation result object. Args: @@ -87,16 +105,20 @@ def _get_aggregation_buckets(self, entry, name=""): yield from self._get_aggregation_buckets(value, name=key) def _run_aggregator( - self, aggregator_name, parameters, search_id=None, chart_type=None - ): + self, + aggregator_name: str, + parameters: Dict[str, Any], + search_id: Optional[int] = None, + chart_type: Optional[str] = None, + ) -> Dict[str, Any]: """Run an aggregator class. Args: - aggregator_name (str): the name of the aggregator class. - parameters (dict): a dict with the parameters for the aggregation class. - search_id (int): an optional integer value with a primary key to a + aggregator_name: the name of the aggregator class. + parameters: a dict with the parameters for the aggregation class. + search_id: an optional integer value with a primary key to a saved search. - chart_type (str): string with the chart type. + chart_type: string with the chart type. Returns: A dict with the aggregation results. @@ -136,11 +158,11 @@ def _run_aggregator( return error.get_response_json(response, logger) # pylint: disable=arguments-renamed - def from_saved(self, aggregation_id): + def from_saved(self, aggregation_id: int) -> None: """Initialize the aggregation object from a saved aggregation. Args: - aggregation_id (int): integer value for the stored + aggregation_id: integer value for the stored aggregation (primary key). """ resource_uri = "sketches/{0:d}/aggregation/{1:d}/".format( @@ -183,11 +205,11 @@ def from_saved(self, aggregation_id): ) # pylint: disable=arguments-differ - def from_manual(self, aggregate_dsl, **kwargs): + def from_manual(self, aggregate_dsl: str, **kwargs: Any) -> None: """Initialize the aggregation object by running an aggregation DSL. Args: - aggregate_dsl (str): OpenSearch aggregation query DSL string. + aggregate_dsl: OpenSearch aggregation query DSL string. kwargs: Optional arguments """ super().from_manual(**kwargs) @@ -217,20 +239,20 @@ def from_manual(self, aggregate_dsl, **kwargs): def from_aggregator_run( self, - aggregator_name, - aggregator_parameters, - search_id=None, - chart_type=None, - ): + aggregator_name: str, + aggregator_parameters: Dict[str, Any], + search_id: Optional[int] = None, + chart_type: Optional[str] = None, + ) -> None: """Initialize the aggregation object by running an aggregator class. Args: - aggregator_name (str): name of the aggregator class to run. - aggregator_parameters (dict): a dict with the parameters of the aggregator + aggregator_name: name of the aggregator class to run. + aggregator_parameters: a dict with the parameters of the aggregator class. - search_id (int): an optional integer value with a primary key to a saved + search_id: an optional integer value with a primary key to a saved search. - chart_type (str): optional string with the chart type. + chart_type: optional string with the chart type. """ self.type = "aggregator_run" self._parameters = aggregator_parameters @@ -244,7 +266,7 @@ def from_aggregator_run( aggregator_name, aggregator_parameters, search_id, chart_type ) - def lazyload_data(self, refresh_cache=False): + def lazyload_data(self, refresh_cache: bool = False) -> Dict[str, Any]: """Load resource data once and cache the result. Args: @@ -260,12 +282,12 @@ def lazyload_data(self, refresh_cache=False): return self.resource_data @property - def parameters(self): + def parameters(self) -> Dict[str, Any]: """Property that returns the parameters of the aggregation.""" return self._parameters @property - def is_part_of_group(self): + def is_part_of_group(self) -> bool: """Property that returns whether an agg is part of a group or not.""" if self._group_id is None: return False @@ -273,7 +295,7 @@ def is_part_of_group(self): return bool(self._group_id) @property - def title(self): + def title(self) -> str: """Property that returns the chart title of an aggregation.""" if self.chart_title: return self.chart_title @@ -287,17 +309,21 @@ def title(self): return self.chart_title @title.setter - def title(self, new_title): - """Set the chart title of an aggregation.""" + def title(self, new_title: str) -> None: + """Set the chart title of an aggregation. + + Args: + new_title: Chart title. + """ self.chart_title = new_title @property - def chart(self): + def chart(self) -> altair.Chart: """Property that returns an altair Vega-lite chart.""" return self.generate_chart() @property - def description(self): + def description(self) -> str: """Property that returns the description string.""" data = self.lazyload_data() if not data: @@ -306,28 +332,37 @@ def description(self): return meta.get("description", "") @description.setter - def description(self, description): - """Set the description of an aggregation.""" - if "meta" not in self.resource_data: + def description(self, description: str) -> None: + """Set the description of an aggregation. + + Args: + description: Description string. + """ + if not self.resource_data or "meta" not in self.resource_data: return meta = self.resource_data.get("meta", {}) meta["description"] = description @property - def name(self): + def name(self) -> str: """Property that returns the name of the aggregation.""" return self._name @name.setter - def name(self, name): - """Set the name of the aggregation.""" - if "meta" not in self.resource_data: + def name(self, name: str) -> None: + """Set the name of the aggregation. + + Args: + name: Name of the aggregation. + """ + if not self.resource_data or "meta" not in self.resource_data: return - meta = self.resource_data.get("meta") + meta = self.resource_data.get("meta") or {} meta["name"] = name + self.resource_data["meta"] = meta @property - def aggregator_name(self): + def aggregator_name(self) -> str: """Property that returns the aggregator name.""" if self._aggregator_name: return self._aggregator_name @@ -338,18 +373,18 @@ def aggregator_name(self): return self._aggregator_name - def add_label(self, label): + def add_label(self, label: str) -> None: """Add a label to the aggregation. Args: - label (str): string with the label information. + label: string with the label information. """ if label in self._labels: return self._labels.append(label) self.save() - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict.""" entries = {} entry_index = 1 @@ -360,7 +395,7 @@ def to_dict(self): entry_index += 1 return entries - def to_pandas(self): + def to_pandas(self) -> pandas.DataFrame: """Returns a pandas DataFrame.""" panda_list = [] data = self.lazyload_data() @@ -370,11 +405,11 @@ def to_pandas(self): return pandas.DataFrame(panda_list) @property - def updated_at(self): + def updated_at(self) -> str: """Returns a timestamp when the aggregation was last updated.""" return self._updated_at - def generate_chart(self): + def generate_chart(self) -> altair.Chart: """Returns an altair Vega-lite chart.""" if not self.chart_type: raise TypeError("Unable to generate chart, missing a chart type.") @@ -392,7 +427,7 @@ def generate_chart(self): vega_spec_string = json.dumps(vega_spec) return altair.Chart.from_json(vega_spec_string) - def save(self): + def save(self) -> str: """Save the aggregation in the database.""" data = { "name": self.name, @@ -425,9 +460,9 @@ def save(self): return "Unable to determine ID of saved object." agg_data = objects[0] self._resource_id = agg_data.get("id", 0) - return "Saved aggregation to ID: {0:d}".format(self._resource_id) + return f"Saved aggregation to ID: {self._resource_id}" - def delete(self): + def delete(self) -> bool: """Deletes the aggregation from the store.""" if not self._resource_id: logger.warning( @@ -446,9 +481,13 @@ def delete(self): class AggregationGroup(resource.SketchResource): """Aggregation Group object.""" - def __init__(self, sketch): - """Initialize the aggregation group.""" - resource_uri = "sketches/{0:d}/aggregation/group/".format(sketch.id) + def __init__(self, sketch: sketch_lib.Sketch) -> None: + """Initialize the aggregation group. + + Args: + sketch: An instance of Sketch object. + """ + resource_uri = f"sketches/{sketch.id}/aggregation/group/" super().__init__(resource_uri=resource_uri, sketch=sketch) self._name = "N/A" @@ -459,84 +498,98 @@ def __init__(self, sketch): self._aggregations = [] self._updated_at = "" - def __str__(self): + def __str__(self) -> str: """Return a string representation of the group.""" - return "[{0:d}] {1:s} - {2:s}".format( - self._resource_id, self._name, self._description - ) + return f"[{self._resource_id}] {self._name} - {self._description}" @property - def aggregations(self): + def aggregations(self) -> List[Aggregation]: """Property that returns a list of aggregations in the group.""" return self._aggregations @property - def created_at(self): + def created_at(self) -> str: """Returns a timestamp when the aggregation group was created.""" return self._created_at @property - def updated_at(self): + def updated_at(self) -> str: """Returns a timestamp when the aggregation group was updated.""" return self._updated_at - def to_dict(self): + def to_dict(self) -> List[Dict[str, Any]]: """Returns the aggregation values as a dict.""" data_frame = self.to_pandas() return data_frame.to_dict(orient="records") @property - def chart(self): + def chart(self) -> altair.Chart: """Property that returns an altair Vega-lite chart.""" if not self._aggregations: return altair.Chart() return self.generate_chart() @property - def description(self): + def description(self) -> str: """Returns the description of the aggregation group.""" return self._description @description.setter - def description(self, description): - """Sets the description of the aggregation group.""" + def description(self, description: str) -> None: + """Sets the description of the aggregation group. + + Args: + description: Description of the aggregation group. + """ self._description = description self.save() @property - def name(self): + def name(self) -> str: """Returns the name of the aggregation group.""" return self._name @name.setter - def name(self, name): - """Sets the name of the aggregation group.""" + def name(self, name: str) -> None: + """Sets the name of the aggregation group. + + Args: + name: Name of the aggregation group. + """ self._name = name self.save() @property - def orientation(self): + def orientation(self) -> str: """Returns the chart orientation.""" return self._orientation @orientation.setter - def orientation(self, orientation): - """Sets the chart orientation.""" + def orientation(self, orientation: str) -> None: + """Sets the chart orientation. + + Args: + orientation: Chart orientation. + """ self._orientation = orientation self.save() @property - def parameters(self): + def parameters(self) -> Dict[str, Any]: """Returns a dict with the group parameters.""" return self._parameters @parameters.setter - def parameters(self, parameters): - """Sets the group parameters.""" + def parameters(self, parameters: Dict[str, Any]) -> None: + """Sets the group parameters. + + Args: + parameters: Aggregation group parameters. + """ self._parameters = parameters self.save() - def delete(self): + def delete(self) -> bool: """Deletes the group from the store.""" if not self._resource_id: logger.warning( @@ -551,11 +604,11 @@ def delete(self): response = self.api.session.delete(resource_uri) return error.check_return_status(response, logger) - def from_dict(self, group_dict): + def from_dict(self, group_dict: Dict[str, Any]) -> None: """Feed group data from a dictionary. Args: - group_dict (dict): a dictionary with the aggregation group + group_dict: a dictionary with the aggregation group information. Raises: @@ -601,11 +654,11 @@ def from_dict(self, group_dict): self._aggregations.append(agg_obj) # pylint: disable=arguments-renamed - def from_saved(self, group_id): + def from_saved(self, group_id: int) -> None: """Feed group data from a group ID. Args: - group_id (int): the group ID to fetch from the store. + group_id: the group ID to fetch from the store. Raises: TypeError: if the group ID does not exist. @@ -622,7 +675,7 @@ def from_saved(self, group_id): group_dict["id"] = group_id self.from_dict(group_dict) - def generate_chart(self): + def generate_chart(self) -> Optional[altair.Chart]: """Returns an altair Vega-lite chart.""" if not self._aggregations: return altair.Chart() @@ -640,15 +693,15 @@ def generate_chart(self): vega_spec_string = json.dumps(vega_spec) return altair.Chart.from_json(vega_spec_string) - def get_charts(self): + def get_charts(self) -> List[altair.Chart]: """Returns a list of altair Chart objects from each aggregation.""" return [x.chart for x in self._aggregations] - def get_tables(self): + def get_tables(self) -> List[pandas.DataFrame]: """Returns a list of pandas DataFrame from each aggregation.""" - return [x.table for x in self._aggregations] + return [x.to_pandas() for x in self._aggregations] - def save(self): + def save(self) -> bool: """Save the aggregation group in the database.""" if not self._aggregations: return False @@ -674,7 +727,7 @@ def save(self): _ = self.lazyload_data(refresh_cache=True) return error.check_return_status(response, logger) - def to_pandas(self): + def to_pandas(self) -> pandas.DataFrame: """Returns a pandas DataFrame. Aggregation groups are meant for charts, not data frames. However diff --git a/api_client/python/timesketch_api_client/aggregation_test.py b/api_client/python/timesketch_api_client/aggregation_test.py index ab3a643ff3..eb23c5c4d7 100644 --- a/api_client/python/timesketch_api_client/aggregation_test.py +++ b/api_client/python/timesketch_api_client/aggregation_test.py @@ -14,7 +14,7 @@ """Tests for the Timesketch API aggregation object.""" import unittest -import mock +from unittest import mock import altair as alt diff --git a/api_client/python/timesketch_api_client/analyzer.py b/api_client/python/timesketch_api_client/analyzer.py index f22690f5e4..bd33c0ca99 100644 --- a/api_client/python/timesketch_api_client/analyzer.py +++ b/api_client/python/timesketch_api_client/analyzer.py @@ -13,23 +13,36 @@ # limitations under the License. """Timesketch API analyzer result object.""" -from __future__ import unicode_literals +from __future__ import annotations import datetime import json import logging +from typing import Any, Dict, Generator, List, TYPE_CHECKING from . import error from . import resource +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.analyzer") class AnalyzerResult(resource.BaseResource): """Class to store and retrieve session information for an analyzer.""" - def __init__(self, timeline_id, session_id, sketch_id, api): - """Initialize the class.""" + def __init__( + self, timeline_id: int, session_id: int, sketch_id: int, api: TimesketchApi + ) -> None: + """Initialize the class. + + Args: + timeline_id: The ID of the timeline. + session_id: The ID of the analyzer session. + sketch_id: The ID of the sketch. + api: An instance of TimesketchApi. + """ self._session_id = session_id self._sketch_id = sketch_id self._timeline_id = timeline_id @@ -38,21 +51,21 @@ def __init__(self, timeline_id, session_id, sketch_id, api): ) super().__init__(api, resource_uri) - def _get_status_data(self): + def _get_status_data(self) -> Generator[Dict[str, str], None, None]: """Yields a dict for each analyzer status.""" data = self._fetch_data() for entry in data.get("analyzers", []): yield { - "log": entry.get("log", "No recorded logs."), - "name": entry.get("name", "No Name"), - "results": entry.get("results", ""), - "status": entry.get("status", "Unknown"), - "date": entry.get( - "status_date", datetime.datetime.utcnow().isoformat() + "log": str(entry.get("log", "No recorded logs.")), + "name": str(entry.get("name", "No Name")), + "results": str(entry.get("results", "")), + "status": str(entry.get("status", "Unknown")), + "date": str( + entry.get("status_date", datetime.datetime.utcnow().isoformat()) ), } - def _fetch_data(self): + def _fetch_data(self) -> Dict[str, Any]: """Returns a dict with the analyzer results.""" response = self.api.session.get(self.resource_uri) if not error.check_return_status(response, logger): @@ -64,7 +77,7 @@ def _fetch_data(self): if not objects: return {} - result_dict = {} + result_dict: Dict[str, Any] = {} for result in objects[0]: result_id = result.get("analysissession_id") if result_id != self._session_id: @@ -100,14 +113,14 @@ def _fetch_data(self): return result_dict @property - def id(self): + def id(self) -> int: """Returns the session ID.""" return self._session_id @property - def log(self): + def log(self) -> str: """Returns back logs from the analyzer session, if there are any.""" - return_strings = [] + return_strings: List[str] = [] for entry in self._get_status_data(): return_strings.append( "[{0:s}] = {1:s}".format( @@ -118,9 +131,9 @@ def log(self): return "\n".join(return_strings) @property - def results(self): + def results(self) -> str: """Returns the results from the analyzer session.""" - return_strings = [] + return_strings: List[str] = [] for entry in self._get_status_data(): results = entry.get("results") if not results: @@ -131,9 +144,9 @@ def results(self): return "\n".join(return_strings) @property - def results_dict(self): + def results_dict(self) -> Dict[str, List[str]]: """Returns the results from the analyzer session as a dict.""" - result_dict = {} + result_dict: Dict[str, List[str]] = {} for entry in self._get_status_data(): results = entry.get("results") if not results: @@ -144,9 +157,9 @@ def results_dict(self): return result_dict @property - def status(self): + def status(self) -> str: """Returns the current status of the analyzer run.""" - return_strings = [] + return_strings: List[str] = [] for entry in self._get_status_data(): return_strings.append( "[{0:s}] = {1:s}".format( @@ -157,9 +170,9 @@ def status(self): return "\n".join(return_strings) @property - def status_dict(self): + def status_dict(self) -> Dict[str, List[str]]: """Returns the current status of the analyzers run as a dict.""" - return_dict = {} + return_dict: Dict[str, List[str]] = {} for entry in self._get_status_data(): name = entry.get("name", "No Name") @@ -168,9 +181,9 @@ def status_dict(self): return return_dict @property - def status_string(self): + def status_string(self) -> str: """Returns a longer version of a status string.""" - return_strings = [] + return_strings: List[str] = [] for entry in self._get_status_data(): return_strings.append( "{0:s} - {1:s}: {2:s}".format( diff --git a/api_client/python/timesketch_api_client/cli_input.py b/api_client/python/timesketch_api_client/cli_input.py index 8e8a7fd024..e577178f17 100644 --- a/api_client/python/timesketch_api_client/cli_input.py +++ b/api_client/python/timesketch_api_client/cli_input.py @@ -13,31 +13,29 @@ # limitations under the License. """CLI assistance for importer tools.""" -from typing import Any -from typing import Callable -from typing import Optional -from typing import Text +from __future__ import annotations import getpass +from typing import Any, Callable, Optional def ask_question( - question: Text, - input_type: Callable[[Text], Any], + question: str, + input_type: Callable[[str], Any], default: Optional[Any] = None, - hide_input: Optional[bool] = False, + hide_input: bool = False, ) -> Any: """Presents the user with a prompt with a default return value and a type. Args: - question (str): the text that the user will be prompted. - input_type (type): the type of the input data. - default (object): default value for the question, optional. - hide_input (bool): whether the input should be hidden, eg. when asking + question: the text that the user will be prompted. + input_type: the type of the input data. + default: default value for the question, optional. + hide_input: whether the input should be hidden, eg. when asking for a password. Returns: - object: The value (type of input_type) that is ready by the user. + The value (type of input_type) that is ready by the user. """ if hide_input: if default: @@ -56,23 +54,21 @@ def ask_question( return input_type(answer) -def confirm_choice( - choice: Text, default: Optional[bool] = True, abort: Optional[bool] = True -) -> bool: +def confirm_choice(choice: str, default: bool = True, abort: bool = True) -> bool: """Returns a bool from a yes/no question presented to the end user. Args: - choice (str): the question presented to the end user. - default (bool): the default for the confirmation answer. If True the + choice: the question presented to the end user. + default: the default for the confirmation answer. If True the default is Y(es), if False the default is N(o) - abort (bool): if the program should abort if the user answer to the + abort: if the program should abort if the user answer to the confirm prompt is no. The default is an abort. Raises: RuntimeError: If abort is set to True and the choice is no. Returns: - bool: False if the user entered no, True if the user entered yes + False if the user entered no, True if the user entered yes """ if default: hint = "Y/n" diff --git a/api_client/python/timesketch_api_client/client.py b/api_client/python/timesketch_api_client/client.py index 621675e0fd..2dfa296a6a 100644 --- a/api_client/python/timesketch_api_client/client.py +++ b/api_client/python/timesketch_api_client/client.py @@ -13,14 +13,15 @@ # limitations under the License. """Timesketch API client.""" -from __future__ import unicode_literals - +from __future__ import annotations import os import logging import sys import time +from typing import Any, Dict, List, Optional, Union, Generator, TYPE_CHECKING + # pylint: disable=wrong-import-order import bs4 import requests @@ -46,6 +47,14 @@ from . import version from . import sigma +if TYPE_CHECKING: + from .credentials import AuthCredentials + from .sketch import Sketch + from .user import User + from .index import SearchIndex + from .sigma import SigmaRule, Sigma + + logger = logging.getLogger("timesketch_api.client") @@ -85,18 +94,18 @@ class TimesketchApi: # pylint: disable=too-many-arguments def __init__( self, - host_uri, - username, - password="", - verify=True, - client_id="", - client_secret="", - auth_mode="userpass", - create_session=True, - retry_count=DEFAULT_RETRY_COUNT, - backoff_factor=0.5, - auth_timeout=None, - ): + host_uri: str, + username: str, + password: str = "", + verify: bool = True, + client_id: str = "", + client_secret: str = "", + auth_mode: str = "userpass", + create_session: bool = True, + retry_count: int = DEFAULT_RETRY_COUNT, + backoff_factor: float = 0.5, + auth_timeout: Optional[int] = None, + ) -> None: """Initializes the TimesketchApi object. Args: @@ -114,7 +123,8 @@ def __init__( function "set_session" needs to be called before proceeding. retry_count: Number of retries for HTTP requests and internal API request retries. Defaults to DEFAULT_RETRY_COUNT. - backoff_factor: The backoff factor to use for retries. Defaults to 0.5. + backoff_factor: The backoff factor to use for retries. + Defaults to 0.5. auth_timeout: Optional timeout in seconds for the authentication. Raises: @@ -157,12 +167,12 @@ def __init__( ) from e @property - def current_user(self): + def current_user(self) -> User: """Property that returns the user object of the logged in user.""" return user.User(self) @property - def version(self): + def version(self) -> str: """Property that returns back the API client version.""" version_dict = self.fetch_resource_data("version/") ts_version = None @@ -177,27 +187,37 @@ def version(self): return "API Client: {0:s}".format(version.get_version()) @property - def session(self): + def session(self) -> requests.Session: """Property that returns the session object.""" if self._session is None: raise ValueError("Session is not set.") return self._session - def set_credentials(self, credential_object): - """Sets the credential object.""" + def set_credentials(self, credential_object: AuthCredentials) -> None: + """Sets the credential object. + + Args: + credential_object: Credential object. + """ self.credentials = credential_object - def set_session(self, session_object): - """Sets the session object.""" + def set_session(self, session_object: requests.Session) -> None: + """Sets the session object. + + Args: + session_object: Instance of requests.Session. + """ self._session = session_object - def _authenticate_session(self, session, username, password): + def _authenticate_session( + self, session: requests.Session, username: str, password: str + ) -> None: """Post username/password to authenticate the HTTP session. Args: - session (requests.Session): Instance of requests.Session. - username (str): User username. - password (str): User password. + session: Instance of requests.Session. + username: User username. + password: User password. """ # Do a POST to the login handler to set up the session cookies data = { @@ -218,12 +238,14 @@ def _authenticate_session(self, session, username, password): if response.url.split("?")[0].rstrip("/").endswith("/login"): raise RuntimeError("Authentication failed: Invalid username or password.") - def _set_csrf_token(self, session, bypass_oauth=False): + def _set_csrf_token( + self, session: requests.Session, bypass_oauth: bool = False + ) -> None: """Retrieve CSRF token from the server and append to HTTP headers. Args: - session (requests.Session): Instance of requests.Session. - bypass_oauth (bool): Whether to bypass OAuth. + session: Instance of requests.Session. + bypass_oauth: Whether to bypass OAuth. """ # Scrape the CSRF token from the response if bypass_oauth: @@ -250,34 +272,34 @@ def _set_csrf_token(self, session, bypass_oauth=False): def _create_oauth_session( self, - client_id="", + client_id: str = "", *, - client_secret="", - client_secrets_file=None, - host="localhost", - port=8080, - open_browser=False, - run_server=True, - skip_open=False, - timeout_seconds=None, - ): + client_secret: str = "", + client_secrets_file: Optional[str] = None, + host: str = "localhost", + port: int = 8080, + open_browser: bool = False, + run_server: bool = True, + skip_open: bool = False, + timeout_seconds: Optional[int] = None, + ) -> requests.Session: """Return an OAuth session. Args: - client_id (str): The client ID if OAUTH auth is used. - client_secret (str): The OAUTH client secret if OAUTH is used. - client_secrets_file (str): Path to the JSON file that contains the client + client_id: The client ID if OAUTH auth is used. + client_secret: The OAUTH client secret if OAUTH is used. + client_secrets_file: Path to the JSON file that contains the client secrets, in the client_secrets format. - host (str): Host address the OAUTH web server will bind to. - port (int): Port the OAUTH web server will bind to. - open_browser (bool): A boolean, if set to false (default) a browser window + host: Host address the OAUTH web server will bind to. + port: Port the OAUTH web server will bind to. + open_browser: A boolean, if set to false (default) a browser window will not be automatically opened. - run_server (bool): A boolean, if set to true (default) a web server is + run_server: A boolean, if set to true (default) a web server is run to catch the OAUTH request and response. - skip_open (bool): A booelan, if set to True (defaults to False) an + skip_open: A booelan, if set to True (defaults to False) an authorization URL is printed on the screen to visit. This is only valid if run_server is set to False. - timeout_seconds (int): Optional timeout in seconds for the authentication. + timeout_seconds: Optional timeout in seconds for the authentication. If set to None (default), it will wait indefinitely. Return: @@ -365,7 +387,7 @@ def _create_oauth_session( self.credentials.credential = flow.credentials return self.authenticate_oauth_session(session) - def authenticate_oauth_session(self, session): + def authenticate_oauth_session(self, session: requests.Session) -> requests.Session: """Authenticate an OAUTH session. Args: @@ -389,31 +411,31 @@ def authenticate_oauth_session(self, session): def _create_session( self, - username, - password, + username: str, + password: str, *, - verify, - client_id, - client_secret, - auth_mode, - retry_count, - backoff_factor, - auth_timeout, - ): + verify: bool, + client_id: str, + client_secret: str, + auth_mode: str, + retry_count: int, + backoff_factor: float, + auth_timeout: Optional[int] = None, + ) -> requests.Session: """Create authenticated HTTP session for server communication. Args: - username (str): User to authenticate as. - password (str): User password. - verify (bool): Verify server SSL certificate. - client_id (str): The client ID if OAUTH auth is used. - client_secret (str): The OAUTH client secret if OAUTH is used. - auth_mode (str): The authentication mode to use. Supported values are + username: User to authenticate as. + password: User password. + verify: Verify server SSL certificate. + client_id: The client ID if OAUTH auth is used. + client_secret: The OAUTH client secret if OAUTH is used. + auth_mode: The authentication mode to use. Supported values are 'userpass' (username/password combo), 'http-basic' (HTTP Basic authentication) and oauth - retry_count (int): Number of retries for HTTP requests. - backoff_factor (float): The backoff factor to use for retries. - auth_timeout (int): Optional timeout in seconds for the authentication. + retry_count: Number of retries for HTTP requests. + backoff_factor: The backoff factor to use for retries. + auth_timeout: Optional timeout in seconds for the authentication. Returns: Instance of requests.Session. @@ -467,7 +489,9 @@ def _create_session( return session - def _send_request_with_retry(self, method, resource_uri, **kwargs): + def _send_request_with_retry( + self, method: str, resource_uri: str, **kwargs: Any + ) -> Dict[str, Any]: """Makes an HTTP request with manual retries for application-level errors. This is the core private helper for all API requests. It wraps the @@ -480,12 +504,12 @@ def _send_request_with_retry(self, method, resource_uri, **kwargs): `GET`. Args: - method (str): HTTP method (e.g., 'GET', 'POST'). - resource_uri (str): The URI for the resource. - **kwargs: Keyword arguments passed to the request. + method: HTTP method (e.g., 'GET', 'POST'). + resource_uri: The URI for the resource. + **kwargs: Optional arguments to send with the request. Returns: - dict: The JSON response data. + The JSON response data. Raises: ValueError: If JSON decoding fails after all retries. @@ -535,7 +559,9 @@ def _send_request_with_retry(self, method, resource_uri, **kwargs): ) raise RuntimeError(error_msg) - def fetch_resource_data(self, resource_uri, params=None): + def fetch_resource_data( + self, resource_uri: str, params: Optional[Dict[str, Any]] = None + ) -> Dict[str, Any]: """Makes an HTTP GET request to the specified resource URI. This method is a convenience wrapper around `_send_request_with_retry` @@ -548,16 +574,16 @@ def fetch_resource_data(self, resource_uri, params=None): JSON responses. Args: - resource_uri (str): The URI to the resource to be fetched. - params (dict, optional): A dictionary of URL parameters to send + resource_uri: The URI to the resource to be fetched. + params: A dictionary of URL parameters to send in the GET request. Defaults to None. Returns: - dict: A dictionary containing the JSON response data from the API. + A dictionary containing the JSON response data from the API. """ return self._send_request_with_retry("GET", resource_uri, params=params) - def create_sketch(self, name, description=None): + def create_sketch(self, name: str, description: Optional[str] = None) -> Sketch: """Create a new sketch. This method attempts to create a new sketch on the Timesketch server. @@ -565,8 +591,8 @@ def create_sketch(self, name, description=None): handle transient network errors or server-side issues. Args: - name (str): Name of the sketch. Cannot be empty. - description (str): Optional description of the sketch. If not + name: Name of the sketch. Cannot be empty. + description: Optional description of the sketch. If not provided, the sketch name will be used as the description. Returns: @@ -611,7 +637,7 @@ def create_sketch(self, name, description=None): ) raise ValueError(error_message_detail) - def create_user(self, username, password): + def create_user(self, username: str, password: str) -> User: """Create a new user. This method attempts to create a new user on the Timesketch server. @@ -619,8 +645,8 @@ def create_user(self, username, password): handle transient network errors or server-side issues. Args: - username (str): Name of the user - password (str): Password of the user + username: Name of the user + password: Password of the user Returns: True if user created successfully. @@ -644,7 +670,7 @@ def create_user(self, username, password): return user.User(user_id=objects[0]["id"], api=self) - def list_users(self): + def list_users(self) -> Generator[User, None, None]: """Get a list of all users. Yields: @@ -657,7 +683,7 @@ def list_users(self): user_obj = user.User(user_id=user_id, api=self) yield user_obj - def get_user(self, user_id): + def get_user(self, user_id: int) -> User: """Get a user. Args: @@ -668,7 +694,7 @@ def get_user(self, user_id): """ return user.User(user_id=user_id, api=self) - def get_oauth_token_status(self): + def get_oauth_token_status(self) -> Dict[str, Any]: """Return a dict with OAuth token status, if one exists.""" if not self.credentials: return {"status": "No stored credentials."} @@ -677,7 +703,7 @@ def get_oauth_token_status(self): "expiry_time": self.credentials.credential.expiry.isoformat(), } - def get_sketch(self, sketch_id): + def get_sketch(self, sketch_id: int) -> Sketch: """Get a sketch. Args: @@ -688,13 +714,15 @@ def get_sketch(self, sketch_id): """ return sketch.Sketch(sketch_id, api=self) - def get_aggregator_info(self, name="", as_pandas=False): + def get_aggregator_info( + self, name: str = "", as_pandas: bool = False + ) -> Union[List[Dict[str, Any]], Dict[str, Any], pandas.DataFrame]: """Returns information about available aggregators. Args: - name (str): String with the name of an aggregator. If the name is not + name: String with the name of an aggregator. If the name is not provided, a list with all aggregators is returned. - as_pandas (bool): Boolean indicating that the results will be returned + as_pandas: Boolean indicating that the results will be returned as a Pandas DataFrame instead of a list of dicts. Returns: @@ -734,12 +762,14 @@ def get_aggregator_info(self, name="", as_pandas=False): return pandas.DataFrame(lines) - def list_sketches(self, per_page=50, scope="user", include_archived=True): + def list_sketches( + self, per_page: int = 50, scope: str = "user", include_archived: bool = True + ) -> Generator[Sketch, None, None]: """Get a list of all open sketches that the user has access to. Args: - per_page (int): Number of items per page when paginating. Default is 50. - scope (str): What scope to get sketches as. Default to user. + per_page: Number of items per page when paginating. Default is 50. + scope: What scope to get sketches as. Default to user. user: sketches owned by the user recent: sketches that the user has actively searched in shared: sketches shared with the user (but not owned by them) @@ -747,7 +777,7 @@ def list_sketches(self, per_page=50, scope="user", include_archived=True): archived: get archived sketches search: pass additional search query all: all sketches the user has access to (owned and shared) - include_archived (bool): If archived sketches should be returned. + include_archived: If archived sketches should be returned. Yields: Sketch objects instances. @@ -778,7 +808,7 @@ def list_sketches(self, per_page=50, scope="user", include_archived=True): ) yield sketch_obj - def get_searchindex(self, searchindex_id): + def get_searchindex(self, searchindex_id: int) -> SearchIndex: """Get a searchindex. Args: @@ -789,7 +819,9 @@ def get_searchindex(self, searchindex_id): """ return index.SearchIndex(searchindex_id, api=self) - def create_searchindex(self, searchindex_name: str, opensearch_index_name: str): + def create_searchindex( + self, searchindex_name: str, opensearch_index_name: str + ) -> SearchIndex: """Create a new SearchIndex. This method attempts to create a new searchindex on the Timesketch server. @@ -837,11 +869,11 @@ def create_searchindex(self, searchindex_name: str, opensearch_index_name: str): ) raise ValueError(error_message_detail) - def check_celery_status(self, job_id=""): + def check_celery_status(self, job_id: str = "") -> List[Dict[str, Any]]: """Return information about outstanding celery tasks or a specific one. Args: - job_id (str): Optional Celery job identification string. If + job_id: Optional Celery job identification string. If provided that specific job ID is queried, otherwise a check for all outstanding jobs is checked. @@ -856,7 +888,7 @@ def check_celery_status(self, job_id=""): return response.get("objects", []) - def list_searchindices(self): + def list_searchindices(self) -> Generator[SearchIndex, None, None]: """Yields all searchindices that the user has access to. Yields: @@ -865,7 +897,6 @@ def list_searchindices(self): response = self.fetch_resource_data("searchindices/") response_objects = response.get("objects") if not response_objects: - yield None return for index_dict in response_objects[0]: @@ -876,20 +907,22 @@ def list_searchindices(self): ) yield index_obj - def refresh_oauth_token(self): + def refresh_oauth_token(self) -> None: """Refresh an OAUTH token if one is defined.""" if not self.credentials: return request = google.auth.transport.requests.Request() self.credentials.credential.refresh(request) - def list_sigmarules(self, as_pandas=False): + def list_sigmarules( + self, as_pandas: bool = False + ) -> Union[List[SigmaRule], pandas.DataFrame]: """Fetches Sigma rules from the database. Fetches all Sigma rules stored in the database on the system and returns a list of SigmaRule objects of the rules. Args: - as_pandas (bool): Boolean indicating that the results will be returned + as_pandas: Boolean indicating that the results will be returned as a Pandas DataFrame instead of a list of SigmaRuleObjects. Returns: @@ -919,7 +952,7 @@ def list_sigmarules(self, as_pandas=False): rules.append(index_obj) return rules - def create_sigmarule(self, rule_yaml): + def create_sigmarule(self, rule_yaml: str) -> SigmaRule: """Adds a single Sigma rule to the database. Adds a single Sigma rule to the database when `/sigmarules/` is called @@ -935,7 +968,7 @@ def create_sigmarule(self, rule_yaml): to handle transient network errors or server-side issues. Args: - rule_yaml (str): YAML of the Sigma Rule. + rule_yaml: YAML of the Sigma Rule. Returns: Instance of a Sigma object. @@ -961,7 +994,7 @@ def create_sigmarule(self, rule_yaml): rule_uuid = objects[0]["rule_uuid"] return self.get_sigmarule(rule_uuid) - def get_sigmarule(self, rule_uuid): + def get_sigmarule(self, rule_uuid: str) -> SigmaRule: """Fetches a single Sigma rule from the database. Fetches a single Sigma rule selected by the `UUID` @@ -976,7 +1009,7 @@ def get_sigmarule(self, rule_uuid): return sigma_obj - def parse_sigmarule_by_text(self, rule_text): + def parse_sigmarule_by_text(self, rule_text: str) -> Sigma: """Obtain a parsed Sigma rule by providing text. Will parse a provided text `rule_yaml`, parse it and return as SigmaRule @@ -1009,12 +1042,16 @@ class VerboseRetry(Retry): MaxRetryError reason. """ - def increment(self, *args, **kwargs): + def increment(self, *args: Any, **kwargs: Any) -> VerboseRetry: """Increment the retry counter and potentially raise MaxRetryError. This method is called by urllib3 before each retry attempt. It's overridden here to add custom logging for 5xx errors and to enhance the final MaxRetryError message with server response details and attempt count. + + Args: + *args: Variable length argument list. + **kwargs: Arbitrary keyword arguments. """ response = kwargs.get("response") decoded_body = None diff --git a/api_client/python/timesketch_api_client/client_test.py b/api_client/python/timesketch_api_client/client_test.py index afcdadd43f..02b80d7e4b 100644 --- a/api_client/python/timesketch_api_client/client_test.py +++ b/api_client/python/timesketch_api_client/client_test.py @@ -13,8 +13,6 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest from unittest import mock diff --git a/api_client/python/timesketch_api_client/config.py b/api_client/python/timesketch_api_client/config.py index b470f06fc6..680a9016d0 100644 --- a/api_client/python/timesketch_api_client/config.py +++ b/api_client/python/timesketch_api_client/config.py @@ -31,6 +31,7 @@ from . import cli_input from . import credentials as ts_credentials from . import crypto +from . import definitions logger = logging.getLogger("timesketch_api.config_assistance") @@ -110,7 +111,7 @@ def get_config(self, name: Text) -> Any: """Returns a value for a given config. Args: - name (str): the name of the config value to retrieve. + name: the name of the config value to retrieve. Raises: KeyError: if the config does not exist. @@ -123,7 +124,7 @@ def get_client( """Returns a Timesketch API client if possible. Args: - token_password (str): an optional password to decrypt + token_password: an optional password to decrypt the credential token file. """ if self.missing: @@ -211,31 +212,31 @@ def has_config(self, name: Text) -> bool: """Returns a boolean indicating whether a config parameter is set. Args: - name (str): the name of the configuration. + name: the name of the configuration. Returns: - bool: whether the object has been set or not. + whether the object has been set or not. """ return name.lower() in self._config def load_config_file( self, config_file_path: Optional[Text] = "", - section: Optional[Text] = "timesketch", + section: Optional[Text] = definitions.DEFAULT_CONFIG_SECTION, load_cli_config: Optional[bool] = False, ): """Load the config from file. Args: - config_file_path (str): Full path to the configuration file, + config_file_path: Full path to the configuration file, if not supplied the default path will be used, which is the file RC_FILENAME inside the user's home directory. - section (str): The configuration section to read from. This + section: The configuration section to read from. This is optional and defaults to timesketch. This can be useful if you have multiple Timesketch servers to connect to, with each one of them having a separate section in the config file. - load_cli_config (bool): Determine if the CLI config section should + load_cli_config: Determine if the CLI config section should be loaded. This is optional and defaults to False. Raises: @@ -271,7 +272,7 @@ def load_config_file( return if not section: - section = "timesketch" + section = definitions.DEFAULT_CONFIG_SECTION if section not in config.sections(): logger.warning("No %s section in the config", section) @@ -304,7 +305,7 @@ def load_config_dict(self, config_dict: Dict[Text, Text]): other keys are ignored in the dict object. Args: - config_dict (dict): dict object with configuration. + config_dict: dict object with configuration. """ fields = list(self.CLIENT_NEEDED) fields.extend(list(self.OAUTH_CLIENT_NEEDED)) @@ -322,21 +323,21 @@ def load_config_dict(self, config_dict: Dict[Text, Text]): def save_config( self, file_path: Optional[Text] = "", - section: Optional[Text] = "timesketch", + section: Optional[Text] = definitions.DEFAULT_CONFIG_SECTION, token_file_path: Optional[Text] = "", ): """Save the current config to a file. Args: - file_path (str): A full path to the location where the + file_path: A full path to the location where the configuration file is to be stored. If not provided the default location will be used. - section (str): The configuration section to write to. This + section: The configuration section to write to. This is optional and defaults to timesketch. This can be useful if you have multiple Timesketch servers to connect to, with each one of them having a separate section in the config file. - token_file_path (str): Optional path to the location of the token + token_file_path: Optional path to the location of the token file. """ if not file_path: @@ -364,7 +365,7 @@ def save_config( auth_mode = "userpass" if not section: - section = "timesketch" + section = definitions.DEFAULT_CONFIG_SECTION config[section] = { "host_uri": self._config.get("host_uri"), @@ -403,8 +404,8 @@ def set_config(self, name: Text, value: Any): """Sets a given config item with a value. Args: - name (str): the name of the configuration value to be set. - value (object): the value of the configuration object. + name: the name of the configuration value to be set. + value: the value of the configuration object. """ self._config[name.lower()] = value @@ -412,7 +413,7 @@ def set_config(self, name: Text, value: Any): def get_client( config_dict: Optional[Dict[Text, Any]] = None, config_path: Optional[Text] = "", - config_section: Optional[Text] = "timesketch", + config_section: Optional[Text] = definitions.DEFAULT_CONFIG_SECTION, token_password: Optional[Text] = "", confirm_choices: Optional[bool] = False, load_cli_config: Optional[bool] = False, @@ -420,21 +421,21 @@ def get_client( """Returns a Timesketch API client using the configuration assistant. Args: - config_dict (dict): optional dict that will be used to configure + config_dict: optional dict that will be used to configure the client. - config_path (str): optional path to the configuration file, if + config_path: optional path to the configuration file, if not supplied a default path will be used. - config_section (str): The configuration section to read from. This + config_section: The configuration section to read from. This is optional and defaults to timesketch. This can be useful if you have multiple Timesketch servers to connect to, with each one of them having a separate section in the config file. - token_password (str): an optional password to decrypt + token_password: an optional password to decrypt the credential token file. - confirm_choices (bool): an optional bool. if set to the user is given + confirm_choices: an optional bool. if set to the user is given a choice to change the value for all already configured parameters. This defaults to False. - load_cli_config (bool): Determine if the CLI config section should + load_cli_config: Determine if the CLI config section should be loaded. This is optional and defaults to False. Returns: @@ -487,7 +488,7 @@ def configure_missing_parameters( config_assistant: ConfigAssistant, token_password: Optional[Text] = "", confirm_choices: Optional[bool] = False, - config_section: Optional[Text] = "timesketch", + config_section: Optional[Text] = definitions.DEFAULT_CONFIG_SECTION, ): """Fill in missing configuration for a config assistant. @@ -499,14 +500,14 @@ def configure_missing_parameters( is username/password and ask for a password to store credentials. Args: - config_assistant (ConfigAssistant): a config assistant that might + config_assistant: a config assistant that might not be fully configured. - token_password (str): an optional password to decrypt + token_password: an optional password to decrypt the credential token file. - confirm_choices (bool): an optional bool. if set to the user is given + confirm_choices: an optional bool. if set to the user is given a choice to change the value for all already configured parameters. This defaults to False. - config_section (str): The configuration section to read from. This + config_section: The configuration section to read from. This is optional and defaults to timesketch. This can be useful if you have multiple Timesketch servers to connect to, with each one of them having a separate section in the config diff --git a/api_client/python/timesketch_api_client/config_test.py b/api_client/python/timesketch_api_client/config_test.py index 1fc9017c05..8756ed3ba8 100644 --- a/api_client/python/timesketch_api_client/config_test.py +++ b/api_client/python/timesketch_api_client/config_test.py @@ -13,8 +13,6 @@ # limitations under the License. """Tests for the Timesketch config library for the API client.""" -from __future__ import unicode_literals - import unittest import tempfile diff --git a/api_client/python/timesketch_api_client/credentials.py b/api_client/python/timesketch_api_client/credentials.py index cbfe997dab..62d9b23401 100644 --- a/api_client/python/timesketch_api_client/credentials.py +++ b/api_client/python/timesketch_api_client/credentials.py @@ -17,9 +17,10 @@ credential objects Timesketch supports. """ -from __future__ import unicode_literals +from __future__ import annotations import json +from typing import Any from google.oauth2 import credentials @@ -30,32 +31,36 @@ class TimesketchCredentials: # The type of credential object. TYPE = "" - def __init__(self): + def __init__(self) -> None: """Initialize the credential object.""" - self._credential = None + self._credential: Any = None @property - def credential(self): + def credential(self) -> Any: """Returns the credentials back.""" return self._credential @credential.setter - def credential(self, credential_obj): - """Sets the credential object.""" + def credential(self, credential_obj: Any) -> None: + """Sets the credential object. + + Args: + credential_obj: The credential object. + """ self._credential = credential_obj - def serialize(self): + def serialize(self) -> bytes: """Return serialized bytes object.""" data = self.to_bytes() type_string = bytes(self.TYPE, "utf-8").rjust(10)[:10] return type_string + data - def deserialize(self, data): + def deserialize(self, data: bytes) -> None: """Deserialize a credential object from bytes. Args: - data (bytes): serialized credential object. + data: serialized credential object. """ type_data = data[:10] type_string = type_data.decode("utf-8").strip() @@ -64,15 +69,15 @@ def deserialize(self, data): self.from_bytes(data[10:]) - def to_bytes(self): + def to_bytes(self) -> bytes: """Convert the credential object into bytes for storage.""" raise NotImplementedError - def from_bytes(self, data): + def from_bytes(self, data: bytes) -> None: """Deserialize a credential object from bytes. Args: - data (bytes): serialized credential object. + data: serialized credential object. """ raise NotImplementedError @@ -82,11 +87,11 @@ class TimesketchPwdCredentials(TimesketchCredentials): TYPE = "timesketch" - def from_bytes(self, data): + def from_bytes(self, data: bytes) -> None: """Deserialize a credential object from bytes. Args: - data (bytes): serialized credential object. + data: serialized credential object. Raises: TypeError: if the data is not in bytes. @@ -99,13 +104,13 @@ def from_bytes(self, data): except ValueError as exc: raise TypeError("Unable to parse the byte string.") from exc - if not "username" in data_dict: + if "username" not in data_dict: raise TypeError("Username is not set.") - if not "password" in data_dict: + if "password" not in data_dict: raise TypeError("Password is not set.") self._credential = data_dict - def to_bytes(self): + def to_bytes(self) -> bytes: """Convert the credential object into bytes for storage.""" if not self._credential: return b"" @@ -119,11 +124,11 @@ class TimesketchOAuthCredentials(TimesketchCredentials): TYPE = "oauth" - def from_bytes(self, data): + def from_bytes(self, data: bytes) -> None: """Deserialize a credential object from bytes. Args: - data (bytes): serialized credential object. + data: serialized credential object. Raises: TypeError: if the data is not in bytes. @@ -145,7 +150,7 @@ def from_bytes(self, data): client_secret=token_dict.get("_client_secret"), ) - def to_bytes(self): + def to_bytes(self) -> bytes: """Convert the credential object into bytes for storage.""" if not self._credential: return b"" diff --git a/api_client/python/timesketch_api_client/crypto.py b/api_client/python/timesketch_api_client/crypto.py index 8423b53e86..87d8b61f0f 100644 --- a/api_client/python/timesketch_api_client/crypto.py +++ b/api_client/python/timesketch_api_client/crypto.py @@ -13,13 +13,14 @@ # limitations under the License. """Timesketch API crypto storage library for OAUTH client.""" -from __future__ import unicode_literals +from __future__ import annotations import base64 import os import getpass import logging import stat +from typing import Any, Optional, Union from cryptography import fernet from cryptography.hazmat import backends @@ -40,8 +41,12 @@ class CredentialStorage: # Length of the salt. SALT_LENGTH = 16 - def __init__(self, file_path=""): - """Initialize the class.""" + def __init__(self, file_path: str = "") -> None: + """Initialize the class. + + Args: + file_path: Path to the credential file. + """ self._user = getpass.getuser() if file_path: self._filepath = file_path @@ -49,12 +54,12 @@ def __init__(self, file_path=""): home_path = os.path.expanduser("~") self._filepath = os.path.join(home_path, self.DEFAULT_CREDENTIAL_FILENAME) - def _get_key(self, salt, password): + def _get_key(self, salt: bytes, password: bytes) -> bytes: """Returns an encryption key. Args: - salt (bytes): a salt used during the encryption. - password (bytes): the password used to decrypt/encrypt + salt: a salt used during the encryption. + password: the password used to decrypt/encrypt the message. Returns: @@ -67,16 +72,24 @@ def _get_key(self, salt, password): iterations=100000, backend=backends.default_backend(), ) - return base64.urlsafe_b64encode(kdf.derive(password)) + return bytes(base64.urlsafe_b64encode(kdf.derive(password))) - def set_filepath(self, file_path): - """Set the filepath to the credential file.""" + def set_filepath(self, file_path: str) -> None: + """Set the filepath to the credential file. + + Args: + file_path: Path to the credential file. + """ if os.path.isfile(file_path): self._filepath = file_path def save_credentials( - self, cred_obj, file_path="", password="", config_assistant=None - ): + self, + cred_obj: credentials.TimesketchCredentials, + file_path: str = "", + password: Union[str, bytes] = "", + config_assistant: Any = None, + ) -> None: """Save credential data to a token file. This function will create an encrypted file @@ -84,14 +97,14 @@ def save_credentials( that contains a stored copy of the credential object. Args: - cred_obj (credentials.TimesketchCredentials): the credential + cred_obj: the credential object that is to be stored on disk. - file_path (str): full path to the file storing the saved + file_path: full path to the file storing the saved credentials. - password (str): optional password to encrypt the + password: optional password to encrypt the credential file with. If not supplied a password will be generated. - config_assistant (ConfigAssistant): optional configuration + config_assistant: optional configuration assistant object. Can be used to store the password to the credential file. """ @@ -99,18 +112,21 @@ def save_credentials( file_path = self._filepath if password: - password = bytes(password, "utf-8") + if isinstance(password, str): + password_bytes = bytes(password, "utf-8") + else: + password_bytes = password else: - password = fernet.Fernet.generate_key() + password_bytes = fernet.Fernet.generate_key() if config_assistant: - config_assistant.set_config("cred_key", password) + config_assistant.set_config("cred_key", password_bytes) config_assistant.save_config() if not os.path.isfile(file_path): logger.info("File does not exist, creating it.") salt = os.urandom(self.SALT_LENGTH) - key = self._get_key(salt, password) + key = self._get_key(salt, password_bytes) crypto = fernet.Fernet(key) data = cred_obj.serialize() @@ -123,16 +139,21 @@ def save_credentials( os.chmod(file_path, file_permission) logger.info("Credentials saved to: %s", file_path) - def load_credentials(self, file_path="", password="", config_assistant=None): + def load_credentials( + self, + file_path: str = "", + password: Union[str, bytes] = "", + config_assistant: Any = None, + ) -> Optional[credentials.TimesketchCredentials]: """Load credentials from a file and return a credential object. Args: - file_path (str): Full path to the file storing the saved + file_path: Full path to the file storing the saved credentials. - password (str): optional password to encrypt the + password: optional password to encrypt the credential file with. If not supplied a password will be generated. - config_assistant (ConfigAssistant): optional configuration + config_assistant: optional configuration assistant object. Can be used to store the password to the credential file. @@ -150,13 +171,19 @@ def load_credentials(self, file_path="", password="", config_assistant=None): if not os.path.isfile(file_path): return None + password_bytes: bytes if password: - password = bytes(password, "utf-8") + if isinstance(password, str): + password_bytes = bytes(password, "utf-8") + else: + password_bytes = password elif config_assistant: try: - password = config_assistant.get_config("cred_key") - if not isinstance(password, bytes): - password = bytes(password, "utf-8") + password_raw = config_assistant.get_config("cred_key") + if not isinstance(password_raw, bytes): + password_bytes = bytes(password_raw, "utf-8") + else: + password_bytes = password_raw except KeyError as exc: raise IOError( "Not able to determine encryption key from config." @@ -169,7 +196,7 @@ def load_credentials(self, file_path="", password="", config_assistant=None): with open(file_path, "rb") as fh: salt = fh.read(self.SALT_LENGTH) - key = self._get_key(salt, password) + key = self._get_key(salt, password_bytes) data = fh.read() crypto = fernet.Fernet(key) try: @@ -185,14 +212,14 @@ def load_credentials(self, file_path="", password="", config_assistant=None): ) from e # TODO: Implement a manager. - cred_obj = credentials.TimesketchPwdCredentials() + pwd_cred_obj = credentials.TimesketchPwdCredentials() try: - cred_obj.deserialize(data_string) + pwd_cred_obj.deserialize(data_string) - return cred_obj + return pwd_cred_obj except TypeError: logger.debug('Credential object is not "timesketch" auth.') - cred_obj = credentials.TimesketchOAuthCredentials() - cred_obj.deserialize(data_string) - return cred_obj + oauth_cred_obj = credentials.TimesketchOAuthCredentials() + oauth_cred_obj.deserialize(data_string) + return oauth_cred_obj diff --git a/api_client/python/timesketch_api_client/definitions.py b/api_client/python/timesketch_api_client/definitions.py index a573566368..557cc7c338 100644 --- a/api_client/python/timesketch_api_client/definitions.py +++ b/api_client/python/timesketch_api_client/definitions.py @@ -25,3 +25,5 @@ # Convenient buckets of return code families HTTP_STATUS_CODE_20X = [HTTP_STATUS_CODE_OK, HTTP_STATUS_CODE_CREATED] + +DEFAULT_CONFIG_SECTION = "timesketch" diff --git a/api_client/python/timesketch_api_client/error.py b/api_client/python/timesketch_api_client/error.py index 85ad8b0989..b3f1cb5930 100644 --- a/api_client/python/timesketch_api_client/error.py +++ b/api_client/python/timesketch_api_client/error.py @@ -13,24 +13,29 @@ # limitations under the License. """Timesketch API client library.""" -from __future__ import unicode_literals +from __future__ import annotations import json +from typing import Any, Dict, Optional, Type, TYPE_CHECKING import bs4 from . import definitions +if TYPE_CHECKING: + import logging + import requests -def _get_message(response): + +def _get_message(response: Optional[requests.Response]) -> str: """Return a formatted message string from the response text. Args: - response (requests.Response): a response object from a HTTP + response: a response object from a HTTP request. Returns: - str: a string with the message field extracted from the + a string with the message field extracted from the response.text. """ if response is None: @@ -41,12 +46,12 @@ def _get_message(response): soup = bs4.BeautifulSoup(response_text_raw, features="html.parser") if soup.p: - return soup.p.string # pytype: disable=attribute-error + return str(soup.p.string) # pytype: disable=attribute-error if isinstance(response_text_raw, bytes): response_text = response_text_raw.decode("utf-8") else: - response_text = response_text_raw + response_text = str(response_text_raw) try: response_dict = json.loads(response_text) @@ -56,18 +61,18 @@ def _get_message(response): if not isinstance(response_dict, dict): return str(response_dict) - return response_dict.get("message", str(response_dict)) + return str(response_dict.get("message", str(response_dict))) -def _get_reason(response): +def _get_reason(response: Optional[requests.Response]) -> str: """Return the reason from a response. Args: - response (requests.Response): a response object from a HTTP + response: a response object from a HTTP request. Returns: - str: a string with the reason field extracted from the + a string with the reason field extracted from the response.reason. """ if response is None: @@ -76,19 +81,21 @@ def _get_reason(response): if isinstance(reason, bytes): return reason.decode("utf-8") - return reason + return str(reason) -def get_response_json(response, logger): +def get_response_json( + response: requests.Response, logger: logging.Logger +) -> Dict[str, Any]: """Return the JSON object from a response, logging any errors. Args: - response (requests.Response): a response object from a HTTP request. - logger (logging.Logger): a logger object that can be used to write log + response: a response object from a HTTP request. + logger: a logger object that can be used to write log messages. Returns: - dict: a dict with the decoded JSON object within the HTTP + a dict with the decoded JSON object within the HTTP response object. Raises: @@ -119,13 +126,17 @@ def get_response_json(response, logger): raise ValueError("Unable to JSON decode the Timesketch API response.") from e -def error_message(response, message=None, error=RuntimeError): +def error_message( + response: requests.Response, + message: Optional[str] = None, + error: Type[Exception] = RuntimeError, +) -> None: """Raise an error using error message extracted from response. Args: - response (requests.Response): a response object from a HTTP request. - message (str): Optional message to prepend to the error string. - error (Exception): The exception class to raise. Defaults to RuntimeError. + response: a response object from a HTTP request. + message: Optional message to prepend to the error string. + error: The exception class to raise. Defaults to RuntimeError. Raises: error: The exception specified by the error argument. @@ -142,29 +153,29 @@ def error_message(response, message=None, error=RuntimeError): ) -def check_return_status(response, logger): +def check_return_status(response: requests.Response, logger: logging.Logger) -> bool: """Check return status and return a boolean. Args: - response (requests.Response): a response object from a HTTP + response: a response object from a HTTP request. - logger (logging.Logger): a logger object that can be used to + logger: a logger object that can be used to write log messages. Returns: - bool: a boolean indicating whether the return status was in + a boolean indicating whether the return status was in the 20X range of HTTP responses. """ status = response.status_code in definitions.HTTP_STATUS_CODE_20X if status: - return status + return bool(status) logger.warning( "Failed response: [{0:d}] {1:s}".format( response.status_code, _get_message(response) ) ) - return status + return bool(status) class Error(Exception): diff --git a/api_client/python/timesketch_api_client/error_test.py b/api_client/python/timesketch_api_client/error_test.py index f133f3f08b..25f671a300 100644 --- a/api_client/python/timesketch_api_client/error_test.py +++ b/api_client/python/timesketch_api_client/error_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client error handling.""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import error diff --git a/api_client/python/timesketch_api_client/graph.py b/api_client/python/timesketch_api_client/graph.py index 88aed085ed..d8e32cf151 100644 --- a/api_client/python/timesketch_api_client/graph.py +++ b/api_client/python/timesketch_api_client/graph.py @@ -49,7 +49,11 @@ class Graph(resource.SketchResource): } def __init__(self, sketch): - """Initialize the graph object.""" + """Initialize the graph object. + + Args: + sketch (Sketch): The sketch that the graph belongs to. + """ resource_uri = f"sketches/{sketch.id}/graphs/" super().__init__(sketch=sketch, resource_uri=resource_uri) @@ -140,7 +144,11 @@ def description(self): @description.setter def description(self, description): - """Make changes to the saved search description field.""" + """Make changes to the saved search description field. + + Args: + description (str): The description of the graph. + """ self._description = description self.commit() @@ -208,7 +216,11 @@ def graph_config(self): @graph_config.setter def graph_config(self, graph_config): - """Change the graph config.""" + """Change the graph config. + + Args: + graph_config (dict): The configuration for the graph. + """ if not isinstance(graph_config, dict): raise ValueError("Graph config needs to be a dict.") @@ -348,7 +360,11 @@ def layout(self): @layout.setter def layout(self, layout): - """Change the layout manually.""" + """Change the layout manually. + + Args: + layout (dict): The layout for the graph. + """ if not isinstance(layout, dict): raise ValueError("Layout needs to be a dict.") self._layout = layout @@ -365,7 +381,11 @@ def name(self): @name.setter def name(self, name): - """Make changes to the saved search name.""" + """Make changes to the saved search name. + + Args: + name (str): The name of the graph. + """ self._name = name self.commit() @@ -427,7 +447,11 @@ def save(self): return f"Saved graph to ID: {self._resource_id}" def set_layout_type(self, layout_string): - """Use a layout from the layout strings.""" + """Use a layout from the layout strings. + + Args: + layout_string (str): The layout to use. + """ layout = self._GRAPH_LAYOUTS.get(layout_string) if layout: self.layout = layout(self.graph) @@ -439,7 +463,11 @@ def timelines(self): @timelines.setter def timelines(self, timelines): - """Sets the timelines.""" + """Sets the timelines. + + Args: + timelines (list): A list of timelines to use. + """ if not isinstance(timelines, (list, tuple)): logger.error("Unable to add timelines, this needs to be a list") diff --git a/api_client/python/timesketch_api_client/graph_test.py b/api_client/python/timesketch_api_client/graph_test.py index de523f18b1..4867039bb6 100644 --- a/api_client/python/timesketch_api_client/graph_test.py +++ b/api_client/python/timesketch_api_client/graph_test.py @@ -14,7 +14,7 @@ """Tests for the Timesketch API client""" import unittest -import mock +from unittest import mock from . import client from . import graph diff --git a/api_client/python/timesketch_api_client/index.py b/api_client/python/timesketch_api_client/index.py index 7a915ea1cd..99f3631c9f 100644 --- a/api_client/python/timesketch_api_client/index.py +++ b/api_client/python/timesketch_api_client/index.py @@ -13,14 +13,23 @@ # limitations under the License. """Timesketch API client library.""" -from __future__ import unicode_literals +from __future__ import annotations import json import logging +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import TYPE_CHECKING + from . import error from . import resource +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.index") @@ -32,20 +41,26 @@ class SearchIndex(resource.BaseResource): api: An instance of TimesketchApi object. """ - def __init__(self, searchindex_id, api, searchindex_name=None): + def __init__( + self, + searchindex_id: int, + api: TimesketchApi, + searchindex_name: Optional[str] = None, + ) -> None: """Initializes the SearchIndex object. Args: searchindex_id: Primary key ID of the searchindex. + api: An instance of TimesketchApi object. searchindex_name: Name of the searchindex (optional). """ self.id = searchindex_id - self._labels = [] + self._labels: List[str] = [] self._searchindex_name = searchindex_name resource_uri = f"searchindices/{self.id}/" super().__init__(api=api, resource_uri=resource_uri) - def _get_object_dict(self): + def _get_object_dict(self) -> Dict[str, Any]: """Returns the object dict from the resources dict.""" data = self.lazyload_data() objects = data.get("objects", []) @@ -55,18 +70,18 @@ def _get_object_dict(self): return objects[0] @property - def fields(self): + def fields(self) -> List[Dict[str, Any]]: """Property that returns the fields in the index mappings.""" index_data = self.lazyload_data(refresh_cache=True) meta = index_data.get("meta", {}) return meta.get("fields", []) @property - def has_timeline_id(self): + def has_timeline_id(self) -> bool: """Property that returns back whether a __ts_timeline_id field is set. Returns: - bool: True if the data uses __timeline_id field to distinguish + True if the data uses __timeline_id field to distinguish different data sets in an index, False if the entire index is the data set. """ @@ -75,7 +90,7 @@ def has_timeline_id(self): return meta.get("contains_timeline_id", False) @property - def labels(self): + def labels(self) -> List[str]: """Property that returns the SearchIndex labels.""" if self._labels: return self._labels @@ -93,7 +108,7 @@ def labels(self): return self._labels @property - def name(self): + def name(self) -> str: """Property that returns searchindex name. Returns: @@ -105,7 +120,7 @@ def name(self): return self._searchindex_name @property - def index_name(self): + def index_name(self) -> str: """Property that returns OpenSearch index name. Returns: @@ -115,7 +130,7 @@ def index_name(self): return index_data.get("index_name", "unknown index name") @property - def status(self): + def status(self) -> str: """Property that returns the index status. Returns: @@ -130,8 +145,12 @@ def status(self): return status.get("status") @status.setter - def status(self, status): - """Set the SearchIndex status.""" + def status(self, status: str) -> None: + """Set the SearchIndex status. + + Args: + status: The status to set. + """ resource_url = f"{self.api.api_root}/searchindices/{self.id}/" data = {"status": status} response = self.api.session.post(resource_url, json=data) @@ -139,12 +158,12 @@ def status(self, status): _ = error.check_return_status(response, logger) @property - def description(self): + def description(self) -> str: """Property that returns the description of the index.""" index_data = self._get_object_dict() return index_data.get("description", "no description provided") - def delete(self): + def delete(self) -> bool: """Deletes the index.""" resource_url = "{0:s}/searchindices/{1:d}/".format(self.api.api_root, self.id) response = self.api.session.delete(resource_url) diff --git a/api_client/python/timesketch_api_client/resource.py b/api_client/python/timesketch_api_client/resource.py index 407918373d..6ba9c18f00 100644 --- a/api_client/python/timesketch_api_client/resource.py +++ b/api_client/python/timesketch_api_client/resource.py @@ -13,7 +13,15 @@ # limitations under the License. """Timesketch API client library.""" +from __future__ import annotations + import json +from typing import Any, Dict, List, Optional, TYPE_CHECKING + +if TYPE_CHECKING: + import pandas + from .client import TimesketchApi + from .sketch import Sketch class BaseResource: @@ -23,7 +31,7 @@ class BaseResource: resources, such as lazy loading of data. """ - def __init__(self, api, resource_uri): + def __init__(self, api: TimesketchApi, resource_uri: str) -> None: """Initializes the BaseResource. Args: @@ -33,9 +41,9 @@ def __init__(self, api, resource_uri): """ self.api = api self.resource_uri = resource_uri - self.resource_data = None + self.resource_data: Optional[Dict[str, Any]] = None - def lazyload_data(self, refresh_cache=False): + def lazyload_data(self, refresh_cache: bool = False) -> Dict[str, Any]: """Load resource data once and cache the result. This method fetches data from the API for this resource if it hasn't @@ -43,21 +51,21 @@ def lazyload_data(self, refresh_cache=False): data is stored in `self.resource_data`. Args: - refresh_cache (bool): If True, forces a refresh of the cached data. + refresh_cache: If True, forces a refresh of the cached data. Returns: - dict: A dictionary containing the resource data from the API. + A dictionary containing the resource data from the API. """ if not self.resource_data or refresh_cache: self.resource_data = self.api.fetch_resource_data(self.resource_uri) return self.resource_data @property - def data(self): + def data(self) -> Dict[str, Any]: """Property to access the resource's data. Returns: - dict: A dictionary containing the resource data from the API. + A dictionary containing the resource data from the API. """ return self.lazyload_data() @@ -65,7 +73,7 @@ def data(self): class SketchResource(BaseResource): """Sketch resource object.""" - def __init__(self, resource_uri, sketch): + def __init__(self, resource_uri: str, sketch: Sketch) -> None: """Initialize the sketch resource object. Args: @@ -74,12 +82,14 @@ def __init__(self, resource_uri, sketch): """ super().__init__(sketch.api, resource_uri) - self._labels = [] + self._labels: List[str] = [] self._resource_id = 0 self._sketch = sketch self._username = "" - def _get_top_level_attribute(self, name, default_value=None, refresh=False): + def _get_top_level_attribute( + self, name: str, default_value: Any = None, refresh: bool = False + ) -> Any: """Returns a top level attribute from a resource object. Args: @@ -91,8 +101,8 @@ def _get_top_level_attribute(self, name, default_value=None, refresh=False): Returns: The dict value of the key "name". """ - resource = self.lazyload_data(refresh_cache=refresh) - resource_objects = resource.get("objects") + resource_data = self.lazyload_data(refresh_cache=refresh) + resource_objects = resource_data.get("objects") if not resource_objects: return default_value @@ -102,40 +112,40 @@ def _get_top_level_attribute(self, name, default_value=None, refresh=False): first_object = resource_objects[0] return first_object.get(name, default_value) - def add_label(self, label): + def add_label(self, label: str) -> None: """Add a label to the resource. Args: - label (str): string with the label information. + label: string with the label information. """ if label in self._labels: return self._labels.append(label) self.save() - def commit(self): + def commit(self) -> None: """Calls the save function if the object has already been saved.""" if not self._resource_id: return self.save() - def delete(self): + def delete(self) -> Any: """Deletes the resource from the list of stored resources.""" raise NotImplementedError @property - def dict(self): + def dict(self) -> Dict[str, Any]: """Property that returns back a Dict with the results.""" return self.to_dict() - def from_manual(self, **kwargs): + def from_manual(self, **kwargs: Any) -> None: """Initialize the resource object by running a raw API request. The API request functionality should be implemented by other functions that inherit this as a base class. Args: - kwargs (dict[str, object]): Depending on the resource they may + kwargs: Depending on the resource they may require different sets of arguments to be able to run a raw API request. @@ -148,7 +158,7 @@ def from_manual(self, **kwargs): "Unused keyword arguments: {0:s}.".format(", ".join(kwargs.keys())) ) - def from_saved(self, resource_id): + def from_saved(self, resource_id: int) -> None: """Initialize the resource object from a saved resource. Args: @@ -157,45 +167,45 @@ def from_saved(self, resource_id): raise NotImplementedError @property - def id(self): + def id(self) -> int: """Property that returns back the resource ID.""" return self._resource_id @property - def json(self): + def json(self) -> str: """Property that returns back a JSON object with the results.""" return json.dumps(self.dict) @property - def labels(self): + def labels(self) -> List[str]: """Property that returns a list of the resource labels.""" return self._labels @property - def sketch(self): + def sketch(self) -> Sketch: """Property that returns the sketch object.""" return self._sketch @property - def table(self): + def table(self) -> pandas.DataFrame: """Property that returns a pandas DataFrame.""" return self.to_pandas() @property - def user(self): + def user(self) -> str: """Property that returns the username of who ran the aggregation.""" if not self._username: return "System" return self._username - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict.""" raise NotImplementedError - def to_pandas(self): + def to_pandas(self) -> pandas.DataFrame: """Returns a pandas DataFrame.""" raise NotImplementedError - def save(self): + def save(self) -> Any: """Sends a request to save the resource.""" raise NotImplementedError diff --git a/api_client/python/timesketch_api_client/scenario.py b/api_client/python/timesketch_api_client/scenario.py index 2a50b64fde..d33a260b77 100644 --- a/api_client/python/timesketch_api_client/scenario.py +++ b/api_client/python/timesketch_api_client/scenario.py @@ -40,8 +40,8 @@ def __init__(self, sketch_id, scenario_id, uuid, api): """Initializes the Scenario object. Args: - sketch_id: ID of a sketch. - scenario_id: Primary key ID of the scenario. + sketch_id (int): ID of a sketch. + scenario_id (int): Primary key ID of the scenario. uuid: UUID of the scenario. api: An instance of the TimesketchApi object. """ @@ -89,11 +89,11 @@ def lazyload_data(self, refresh_cache: bool = False) -> Dict[str, Any]: True. Args: - refresh_cache (bool): If True, forces a refresh of the cached data + refresh_cache: If True, forces a refresh of the cached data from the API. Returns: - dict: A dictionary containing the resource data from the API. + A dictionary containing the resource data from the API. """ if not self._is_populated or refresh_cache: super().lazyload_data(refresh_cache=refresh_cache) @@ -150,7 +150,7 @@ def list_facets(self) -> List[Dict[str, Any]]: """Lists all facets for the scenario. Returns: - list[dict]: A list of dictionaries, each representing a facet. + A list of dictionaries, each representing a facet. """ resource_url = ( f"{self.api.api_root}/sketches/{self.sketch_id}/" @@ -206,8 +206,8 @@ def __init__(self, sketch_id, question_id, uuid, api): """Initializes the Question object. Args: - sketch_id: ID of a sketch. - question_id: Primary key ID of the question. + sketch_id (int): ID of a sketch. + question_id (int): Primary key ID of the question. uuid: UUID of the question. api: An instance of the TimesketchApi object. """ @@ -257,11 +257,11 @@ def lazyload_data(self, refresh_cache: bool = False) -> Dict[str, Any]: True. Args: - refresh_cache (bool): If True, forces a refresh of the cached data + refresh_cache: If True, forces a refresh of the cached data from the API. Returns: - dict: A dictionary containing the resource data from the API. + A dictionary containing the resource data from the API. """ if not self._is_populated or refresh_cache: super().lazyload_data(refresh_cache=refresh_cache) diff --git a/api_client/python/timesketch_api_client/scenario_test.py b/api_client/python/timesketch_api_client/scenario_test.py index 23f6e7ae81..8bbd489b39 100644 --- a/api_client/python/timesketch_api_client/scenario_test.py +++ b/api_client/python/timesketch_api_client/scenario_test.py @@ -14,7 +14,7 @@ """Tests for the Timesketch API client""" import unittest -import mock +from unittest import mock from . import client from . import test_lib diff --git a/api_client/python/timesketch_api_client/search.py b/api_client/python/timesketch_api_client/search.py index 5d27d0cd50..17a7d165ae 100644 --- a/api_client/python/timesketch_api_client/search.py +++ b/api_client/python/timesketch_api_client/search.py @@ -13,17 +13,29 @@ # limitations under the License. """Timesketch API search object.""" +from __future__ import annotations + import datetime import json import logging import re +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import Union +from typing import TYPE_CHECKING + import pandas from . import error from . import resource from . import searchtemplate +if TYPE_CHECKING: + from . import sketch as sketch_lib + logger = logging.getLogger("timesketch_api.search") @@ -40,24 +52,28 @@ class Chip: # The value of the chip field. CHIP_FIELD = "" - def __init__(self): + def __init__(self) -> None: """Initialize the chip.""" self._active = True self._operator = "must" self._chip_field = self.CHIP_FIELD @property - def active(self): + def active(self) -> bool: """A property that returns whether the chip is active or not.""" return self._active @active.setter - def active(self, active): - """Decide whether the chip is active or disabled.""" + def active(self, active: bool) -> None: + """Decide whether the chip is active or disabled. + + Args: + active: Boolean indicating if the chip is active. + """ self._active = bool(active) @property - def chip(self): + def chip(self) -> Dict[str, Any]: """A property that returns the chip value.""" return { "field": self._chip_field, @@ -67,27 +83,31 @@ def chip(self): "value": getattr(self, self.CHIP_VALUE, ""), } - def from_dict(self, chip_dict): - """Configure the chip from a dictionary.""" + def from_dict(self, chip_dict: Dict[str, Any]) -> None: + """Configure the chip from a dictionary. + + Args: + chip_dict: Dictionary with chip configuration. + """ raise NotImplementedError - def set_include(self): + def set_include(self) -> None: """Configure the chip so the content needs to be included in results.""" self._operator = "must" - def set_exclude(self): + def set_exclude(self) -> None: """Configure the chip so content needs to be excluded in results.""" self._operator = "must_not" - def set_optional(self): + def set_optional(self) -> None: """Configure the chip so the content is optional in results.""" self._operator = "should" - def set_active(self): + def set_active(self) -> None: """Set the chip as active.""" self._active = True - def set_disable(self): + def set_disable(self) -> None: """Disable the chip.""" self._active = False @@ -102,7 +122,7 @@ class DateIntervalChip(Chip): _DATE_FORMAT_MICROSECONDS = "%Y-%m-%dT%H:%M:%S.%f" _DATE_ONLY_FORMAT = "%Y-%m-%d" - def __init__(self): + def __init__(self) -> None: """Initialize the chip.""" super().__init__() self._date = None @@ -110,15 +130,17 @@ def __init__(self): self._after = 5 self._unit = "m" - def add_interval(self, before, after=None, unit="m"): + def add_interval( + self, before: int, after: Optional[int] = None, unit: str = "m" + ) -> None: """Set the interval of the chip. Args: - before (int): the number of units that should be included + before: the number of units that should be included before the date. - after (int): optional number of units after the date. If not + after: optional number of units after the date. If not provided the value of before is used. - unit (str): optional string with the unit of interval. This can + unit: optional string with the unit of interval. This can be s for seconds, m for minutes, d for days and h for hours. The default value is m (minutes). @@ -134,27 +156,35 @@ def add_interval(self, before, after=None, unit="m"): self._after = after @property - def after(self): + def after(self) -> int: """Property that returns the time interval after the date.""" return self._after @after.setter - def after(self, after): - """Make changes to the time interval after the date.""" + def after(self, after: int) -> None: + """Make changes to the time interval after the date. + + Args: + after: Time interval after the date. + """ self._after = after @property - def before(self): + def before(self) -> int: """Property that returns the time interval before the date.""" return self._before @before.setter - def before(self, before): - """Make changes to the time interval before the date.""" + def before(self, before: int) -> None: + """Make changes to the time interval before the date. + + Args: + before: Time interval before the date. + """ self._before = before @property - def date(self): + def date(self) -> str: """Property that returns back the date.""" if not self._date: return "" @@ -163,8 +193,12 @@ def date(self): return self._date.strftime(self._DATE_FORMAT_MICROSECONDS)[:-3] @date.setter - def date(self, date): - """Make changes to the date.""" + def date(self, date: str) -> None: + """Make changes to the date. + + Args: + date: Date string. + """ try: dt = datetime.datetime.strptime(date, self._DATE_FORMAT_MICROSECONDS) except ValueError: @@ -180,8 +214,12 @@ def date(self, date): raise ValueError("Wrong date format") from exc self._date = dt - def from_dict(self, chip_dict): - """Configure the chip from a dictionary.""" + def from_dict(self, chip_dict: Dict[str, Any]) -> None: + """Configure the chip from a dictionary. + + Args: + chip_dict: Dictionary with chip configuration. + """ value = chip_dict.get("value") if not value: return @@ -204,18 +242,22 @@ def from_dict(self, chip_dict): self.after = int(after[1:-1]) @property - def interval(self): + def interval(self) -> str: """A property that returns back the full interval.""" return f"{self.date} -{self.before}{self.unit} +{self.after}{self.unit}" @property - def unit(self): + def unit(self) -> str: """Property that returns back the unit used.""" return self._unit @unit.setter - def unit(self, unit): - """Make changes to the unit.""" + def unit(self, unit: str) -> None: + """Make changes to the unit. + + Args: + unit: Unit of interval. + """ if unit not in ("s", "m", "d", "h"): raise ValueError( "Unable to add interval, needs to be one of: " @@ -235,18 +277,18 @@ class DateRangeChip(Chip): _DATE_RE = r"^[0-9]{4}-[0-9]{1,2}-[0-9]{2}$" - def __init__(self): + def __init__(self) -> None: """Initialize the date range.""" super().__init__() self._start_date = None self._end_date = None self._date_re = re.compile(self._DATE_RE) - def add_end_time(self, end_time): + def add_end_time(self, end_time: str) -> None: """Add an end time to the range. Args: - end_time (str): date string using the format '%Y-%m-%dT%H:%M:%s' + end_time: date string using the format '%Y-%m-%dT%H:%M:%s' Raises: ValueError: if the date format is incorrectly formatted. @@ -270,11 +312,11 @@ def add_end_time(self, end_time): raise ValueError("Wrong date format") from exc self._end_date = dt - def add_start_time(self, start_time): + def add_start_time(self, start_time: str) -> None: """Add a start time to the range. Args: - start_time (str): date string using the format '%Y-%m-%dT%H:%M:%s' + start_time: date string using the format '%Y-%m-%dT%H:%M:%s' Raises: ValueError: if the date format is incorrectly formatted. @@ -299,7 +341,7 @@ def add_start_time(self, start_time): self._start_date = dt @property - def end_time(self): + def end_time(self) -> str: """Property that returns the end time of a range.""" if not self._end_date: return "" @@ -308,24 +350,36 @@ def end_time(self): return self._end_date.strftime(self._DATE_FORMAT_MICROSECONDS)[:-3] @end_time.setter - def end_time(self, end_time): - """Sets the new end time.""" + def end_time(self, end_time: str) -> None: + """Sets the new end time. + + Args: + end_time: New end time. + """ self.add_end_time(end_time) @property - def date_range(self): + def date_range(self) -> str: """Property that returns back the range.""" return f"{self.start_time},{self.end_time}" @date_range.setter - def date_range(self, date_range): - """Sets the new range of the date range chip.""" + def date_range(self, date_range: str) -> None: + """Sets the new range of the date range chip. + + Args: + date_range: Comma separated string with start and end time. + """ start_time, end_time = date_range.split(",") self.add_start_time(start_time) self.add_end_time(end_time) - def from_dict(self, chip_dict): - """Configure the chip from a dictionary.""" + def from_dict(self, chip_dict: Dict[str, Any]) -> None: + """Configure the chip from a dictionary. + + Args: + chip_dict: Dictionary with chip configuration. + """ chip_value = chip_dict.get("value") if not chip_value: return @@ -334,7 +388,7 @@ def from_dict(self, chip_dict): self.end_time = end @property - def start_time(self): + def start_time(self) -> str: """Property that returns the start time of a range.""" if not self._start_date: return "" @@ -343,8 +397,12 @@ def start_time(self): return self._start_date.strftime(self._DATE_FORMAT_MICROSECONDS)[:-3] @start_time.setter - def start_time(self, start_time): - """Sets the new start time of a range.""" + def start_time(self, start_time: str) -> None: + """Sets the new start time of a range. + + Args: + start_time: New start time. + """ self.add_start_time(start_time) @@ -354,13 +412,17 @@ class LabelChip(Chip): CHIP_TYPE = "label" CHIP_VALUE = "label" - def __init__(self): + def __init__(self) -> None: """Initialize the chip.""" super().__init__() self._label = "" - def from_dict(self, chip_dict): - """Configure the chip from a dictionary.""" + def from_dict(self, chip_dict: Dict[str, Any]) -> None: + """Configure the chip from a dictionary. + + Args: + chip_dict: Dictionary with chip configuration. + """ chip_value = chip_dict.get("value") if not chip_value: return @@ -368,20 +430,24 @@ def from_dict(self, chip_dict): self.label = chip_value @property - def label(self): + def label(self) -> str: """Property that returns back the label.""" return self._label @label.setter - def label(self, label): - """Make changes to the label.""" + def label(self, label: str) -> None: + """Make changes to the label. + + Args: + label: Label string. + """ self._label = label - def use_comment_label(self): + def use_comment_label(self) -> None: """Use the comment label.""" self._label = "__ts_comment" - def use_star_label(self): + def use_star_label(self) -> None: """Use the star label.""" self._label = "__ts_star" @@ -392,23 +458,31 @@ class TermChip(Chip): CHIP_TYPE = "term" CHIP_VALUE = "query" - def __init__(self): + def __init__(self) -> None: """Initialize the chip.""" super().__init__() self._query = "" @property - def field(self): + def field(self) -> str: """Property that returns back the field used to match against.""" return self._chip_field @field.setter - def field(self, field): - """Make changes to the field used to match against.""" + def field(self, field: str) -> None: + """Make changes to the field used to match against. + + Args: + field: Field name. + """ self._chip_field = field - def from_dict(self, chip_dict): - """Configure the term chip from a dictionary.""" + def from_dict(self, chip_dict: Dict[str, Any]) -> None: + """Configure the term chip from a dictionary. + + Args: + chip_dict: Dictionary with chip configuration. + """ chip_value = chip_dict.get("value") if not chip_value: return @@ -417,13 +491,17 @@ def from_dict(self, chip_dict): self.query = chip_value @property - def query(self): + def query(self) -> str: """Property that returns back the query.""" return self._query @query.setter - def query(self, query): - """Make changes to the query.""" + def query(self, query: str) -> None: + """Make changes to the query. + + Args: + query: Query string. + """ self._query = query @@ -432,7 +510,12 @@ class Search(resource.SketchResource): DEFAULT_SIZE_LIMIT = 10000 - def __init__(self, sketch): + def __init__(self, sketch: sketch_lib.Sketch) -> None: + """Initialize the Search object. + + Args: + sketch: An instance of Sketch object. + """ resource_uri = f"sketches/{sketch.id}/explore/" super().__init__(sketch=sketch, resource_uri=resource_uri) @@ -454,8 +537,12 @@ def __init__(self, sketch): self._updated_at = "" self._use_wildcard_fields = False - def _extract_chips(self, query_filter): - """Extract chips from a query_filter.""" + def _extract_chips(self, query_filter: Dict[str, Any]) -> None: + """Extract chips from a query_filter. + + Args: + query_filter: Query filter dictionary. + """ self._chips = [] chips = query_filter.get("chips", []) if not chips: @@ -491,20 +578,22 @@ def _extract_chips(self, query_filter): self._chips.append(chip) - def _execute_query(self, file_name="", count=False, stream=False): + def _execute_query( + self, file_name: str = "", count: bool = False, stream: bool = False + ) -> Optional[Union[Dict[str, Any], int]]: """Execute a search request and store the results. Args: - file_name (str): Optional file path to a filename that + file_name: Optional file path to a filename that all the results will be saved to. If not provided the results will be stored in the search object. - count (bool): Optional boolean that determines whether + count: Optional boolean that determines whether we want to execute the query or only count the number of events that the query would produce. If set to True, the results will be stored in the search object, and the number of events will be returned. - stream (bool): Optional boolean that determines whether + stream: Optional boolean that determines whether we want to stream the results to a file. This is useful for large exports. @@ -599,18 +688,22 @@ def _execute_query(self, file_name="", count=False, stream=False): self._raw_response = response_json return response_json - def add_chip(self, chip): - """Add a chip to the ...""" + def add_chip(self, chip: Chip) -> None: + """Add a chip to the search object. + + Args: + chip: A chip object. + """ self._chips.append(chip) self.commit() - def add_date_range(self, start_time, end_time): + def add_date_range(self, start_time: str, end_time: str) -> None: """Add a date range chip to the search query. Args: - start_time (str): a string with the start time of the range, + start_time: a string with the start time of the range, the format should be '%Y-%m-%dT%H:%M:%S' - end_time (str): a string with the end time of the range, + end_time: a string with the end time of the range, the format should be '%Y-%m-%dT%H:%M:%S' """ chip = DateRangeChip() @@ -619,21 +712,21 @@ def add_date_range(self, start_time, end_time): self.add_chip(chip) @property - def chips(self): + def chips(self) -> List[Chip]: """Property that returns all the chips in the search object.""" return self._chips - def commit(self): + def commit(self) -> None: """Commit changes to the search object.""" self._raw_response = None super().commit() @property - def created_at(self): + def created_at(self) -> str: """Property that returns back the creation time of a search.""" return self._created_at - def delete(self): + def delete(self) -> bool: """Deletes the saved search from the store.""" if not self._resource_id: logger.warning( @@ -650,18 +743,22 @@ def delete(self): return error.check_return_status(response, logger) @property - def description(self): + def description(self) -> str: """Property that returns back the description of the saved search.""" return self._description @description.setter - def description(self, description): - """Make changes to the saved search description field.""" + def description(self, description: str) -> None: + """Make changes to the saved search description field. + + Args: + description: Description of the saved search. + """ self._description = description self.commit() @property - def expected_size(self): + def expected_size(self) -> int: """Property that returns the expected size of the search query.""" if self._total_elastic_size: return self._total_elastic_size @@ -671,27 +768,27 @@ def expected_size(self): def from_manual( # pylint: disable=arguments-differ self, - query_string=None, - query_dsl=None, - query_filter=None, - return_fields=None, - max_entries=None, - **kwargs, - ): + query_string: Optional[str] = None, + query_dsl: Optional[str] = None, + query_filter: Optional[Dict[str, Any]] = None, + return_fields: Optional[str] = None, + max_entries: Optional[int] = None, + **kwargs: Any, + ) -> None: """Explore the sketch. Args: - query_string (str): OpenSearch query string. - query_dsl (str): OpenSearch query DSL as JSON string. - query_filter (dict): Filter for the query as a dict. - return_fields (str): A comma separated string with a list of fields + query_string: OpenSearch query string. + query_dsl: OpenSearch query DSL as JSON string. + query_filter: Filter for the query as a dict. + return_fields: A comma separated string with a list of fields that should be included in the response. Optional and defaults to None. - max_entries (int): Optional integer denoting a best effort to limit + max_entries: Optional integer denoting a best effort to limit the output size to the number of events. Events are read in, 10k at a time so there may be more events in the answer back than this number denotes, this is a best effort. - kwargs (dict[str, object]): Depending on the resource they may + kwargs: Depending on the resource they may require different sets of arguments to be able to run a raw API request. @@ -727,11 +824,11 @@ def from_manual( # pylint: disable=arguments-differ self.resource_data = {} - def from_saved(self, search_id): # pylint: disable=arguments-renamed + def from_saved(self, search_id: int) -> None: # pylint: disable=arguments-renamed """Initialize the search object from a saved search. Args: - search_id (int): integer value for the saved + search_id: integer value for the saved search (primary key). """ resource_uri = f"sketches/{self._sketch.id}/views/{search_id}/" @@ -775,19 +872,31 @@ def from_saved(self, search_id): # pylint: disable=arguments-renamed self.resource_data = data @property - def indices(self): + def indices(self) -> Union[str, List[str]]: """Return the current set of indices used in the search.""" return self._indices @indices.setter - def indices(self, indices): - """Make changes to the current set of indices.""" + def indices(self, indices: Union[str, List[Union[str, int]]]) -> None: + """Make changes to the current set of indices. + + Args: + indices: List of indices to search. + """ if indices == "_all": self._indices = "_all" self.commit() return - def _is_string_or_int(item): + def _is_string_or_int(item: Any) -> bool: + """Returns whether the item is a string or an int. + + Args: + item: The item to check. + + Returns: + True if the item is a string or an int, False otherwise. + """ return isinstance(item, (str, int)) if not isinstance(indices, list): @@ -853,13 +962,17 @@ def _is_string_or_int(item): self.commit() @property - def max_entries(self): + def max_entries(self) -> int: """Return the maximum number of entries in the return value.""" return self._max_entries @max_entries.setter - def max_entries(self, max_entries): - """Make changes to the max entries of return values.""" + def max_entries(self, max_entries: int) -> None: + """Make changes to the max entries of return values. + + Args: + max_entries: The maximum number of entries to return. + """ self._max_entries = max_entries if max_entries < self.DEFAULT_SIZE_LIMIT: _ = self.query_filter @@ -868,24 +981,28 @@ def max_entries(self, max_entries): self.commit() @property - def name(self): + def name(self) -> str: """Property that returns the query name.""" return self._name @name.setter - def name(self, name): - """Make changes to the saved search name.""" + def name(self, name: str) -> None: + """Make changes to the saved search name. + + Args: + name: Name of the saved search. + """ self._name = name self.commit() - def order_ascending(self): + def order_ascending(self) -> None: """Set the order of objects returned back ascending.""" # Trigger a creation of a query filter if it does not exist. _ = self.query_filter self._query_filter["order"] = "asc" self.commit() - def order_descending(self): + def order_descending(self) -> None: """Set the order of objects returned back descending.""" # Trigger a creation of a query filter if it does not exist. _ = self.query_filter @@ -893,13 +1010,17 @@ def order_descending(self): self.commit() @property - def query_dsl(self): + def query_dsl(self) -> str: """Property that returns back the query DSL.""" return self._query_dsl @query_dsl.setter - def query_dsl(self, query_dsl): - """Make changes to the query DSL of the search.""" + def query_dsl(self, query_dsl: str) -> None: + """Make changes to the query DSL of the search. + + Args: + query_dsl: OpenSearch query DSL as JSON string. + """ if query_dsl and isinstance(query_dsl, str): query_dsl = json.loads(query_dsl) @@ -911,7 +1032,7 @@ def query_dsl(self, query_dsl): self.commit() @property - def query_filter(self): + def query_filter(self) -> Dict[str, Any]: """Property that returns the query filter.""" if not self._query_filter: self._query_filter = { @@ -930,8 +1051,12 @@ def query_filter(self): return query_filter @query_filter.setter - def query_filter(self, query_filter): - """Make changes to the query filter.""" + def query_filter(self, query_filter: Union[str, Dict[str, Any]]) -> None: + """Make changes to the query filter. + + Args: + query_filter: Filter for the query as a dict. + """ if isinstance(query_filter, str): try: query_filter = json.loads(query_filter) @@ -946,13 +1071,18 @@ def query_filter(self, query_filter): self.commit() @property - def use_wildcard_fields(self): + def use_wildcard_fields(self) -> bool: """Return whether wildcard fields search mode is enabled.""" return self._use_wildcard_fields @use_wildcard_fields.setter def use_wildcard_fields(self, enabled: bool) -> None: - """Enable or disable wildcard fields search mode. Defaults to False.""" + """Enable or disable wildcard fields search mode. Defaults to False. + + Args: + enabled: Boolean indicating if wildcard fields search + mode should be enabled. + """ self._use_wildcard_fields = bool(enabled) # Sync to query_filter dictionary immediately on change _ = self.query_filter @@ -960,18 +1090,26 @@ def use_wildcard_fields(self, enabled: bool) -> None: self.commit() @property - def query_string(self): + def query_string(self) -> str: """Property that returns back the query string.""" return self._query_string @query_string.setter - def query_string(self, query_string): - """Make changes to the query string of a saved search.""" + def query_string(self, query_string: str) -> None: + """Make changes to the query string of a saved search. + + Args: + query_string: OpenSearch query string. + """ self._query_string = query_string self.commit() - def remove_chip(self, chip_index): - """Remove a chip from the saved search.""" + def remove_chip(self, chip_index: int) -> None: + """Remove a chip from the saved search. + + Args: + chip_index: Index of the chip to remove. + """ chip_len = len(self._chips) if chip_index > (chip_len + 1): raise ValueError( @@ -989,7 +1127,7 @@ def remove_chip(self, chip_index): self.commit() @property - def return_fields(self): + def return_fields(self) -> str: """Property that returns the return_fields.""" if self._return_fields: items = self._return_fields.split(",") @@ -999,23 +1137,32 @@ def return_fields(self): return self._return_fields @return_fields.setter - def return_fields(self, return_fields): - """Make changes to the return fields.""" + def return_fields(self, return_fields: str) -> None: + """Make changes to the return fields. + + Args: + return_fields: A comma separated string with a list of fields + that should be included in the response. + """ self._return_fields = return_fields self.commit() @property - def return_size(self): + def return_size(self) -> int: """Return the maximum number of entries in the return value.""" return self._max_entries @return_size.setter - def return_size(self, return_size): - """Make changes to the maximum number of entries in the return.""" + def return_size(self, return_size: int) -> None: + """Make changes to the maximum number of entries in the return. + + Args: + return_size: Maximum number of entries to return. + """ self._max_entries = return_size self.commit() - def save(self): + def save(self) -> str: """Save the search in the database. Returns: @@ -1077,7 +1224,7 @@ def save(self): self._resource_id = search_dict.get("id", 0) return f"Saved search to ID: {self._resource_id}" - def save_as_template(self): + def save_as_template(self) -> searchtemplate.SearchTemplate: """Save the search as a search template. Returns: @@ -1096,19 +1243,19 @@ def save_as_template(self): return template @property - def scrolling(self): + def scrolling(self) -> bool: """Returns whether scrolling is enabled or not.""" return self._scrolling - def scrolling_disable(self): + def scrolling_disable(self) -> None: """ "Disables scrolling.""" self._scrolling = False - def scrolling_enable(self): + def scrolling_enable(self) -> None: """Enable scrolling.""" self._scrolling = True - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the response of the query.""" if self._raw_response is None: self._execute_query() @@ -1117,14 +1264,14 @@ def to_dict(self): return self._raw_response - def to_file(self, file_name, stream=False): + def to_file(self, file_name: str, stream: bool = False) -> bool: """Saves the content of the query to a file. Args: - file_name (str): Full path to a file that will store the results + file_name: Full path to a file that will store the results of the query to as a ZIP file. The ZIP file will contain a METADATA file and a CSV with the results from the query. - stream (bool): Optional boolean that determines whether + stream: Optional boolean that determines whether we want to stream the results to a file. This is useful for large exports. @@ -1137,7 +1284,7 @@ def to_file(self, file_name, stream=False): self._scrolling = old_scrolling return True - def to_pandas(self): + def to_pandas(self) -> pandas.DataFrame: """Returns a pandas DataFrame with the response of the query.""" if self._raw_response is None: self._raw_response = self._execute_query() @@ -1190,6 +1337,6 @@ def to_pandas(self): return data_frame @property - def updated_at(self): + def updated_at(self) -> str: """Property that returns back the updated time of a search.""" return self._updated_at diff --git a/api_client/python/timesketch_api_client/search_test.py b/api_client/python/timesketch_api_client/search_test.py index 8a3549031b..6b2dc5b7a6 100644 --- a/api_client/python/timesketch_api_client/search_test.py +++ b/api_client/python/timesketch_api_client/search_test.py @@ -14,7 +14,7 @@ """Tests for the Timesketch API client""" import unittest -import mock +from unittest import mock from . import client from . import search diff --git a/api_client/python/timesketch_api_client/searchtemplate.py b/api_client/python/timesketch_api_client/searchtemplate.py index b30b93c489..01ed52b074 100644 --- a/api_client/python/timesketch_api_client/searchtemplate.py +++ b/api_client/python/timesketch_api_client/searchtemplate.py @@ -26,7 +26,11 @@ class SearchTemplate(resource.BaseResource): """Search template object. TEST e2e""" def __init__(self, api): - """Initialize the search template object.""" + """Initialize the search template object. + + Args: + api (TimesketchApi): An instance of TimesketchApi object. + """ super().__init__(api, "searchtemplates/") self._description = "" self._name = "" diff --git a/api_client/python/timesketch_api_client/sigma.py b/api_client/python/timesketch_api_client/sigma.py index 9cf7055a04..1412706bc6 100644 --- a/api_client/python/timesketch_api_client/sigma.py +++ b/api_client/python/timesketch_api_client/sigma.py @@ -13,13 +13,17 @@ # limitations under the License. """Timesketch API sigma library.""" -from __future__ import unicode_literals +from __future__ import annotations import logging +from typing import Any, Dict, List, TYPE_CHECKING from . import resource from . import error +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.sigma") @@ -32,119 +36,124 @@ class SigmaRule(resource.BaseResource): rule_uuid: The ID of the rule. """ - def __init__(self, api): + def __init__(self, api: TimesketchApi) -> None: """Initializes the Sigma object. Args: api: An instance of TimesketchApi object. - """ - self._attr_dict = {} + self._attr_dict: Dict[str, Any] = {} resource_uri = "sigmarules/" super().__init__(api=api, resource_uri=resource_uri) @property - def attributes(self): + def attributes(self) -> List[str]: """Returns a list of all attribute keys for the rule""" return list(self._attr_dict.keys()) - def get_attribute(self, key): - """Get a value for a given key in case it has no dedicated property""" + def get_attribute(self, key: str) -> Any: + """Get a value for a given key in case it has no dedicated property. + + Args: + key: Key of the attribute to get. + """ if not self._attr_dict: return "" return self._attr_dict.get(key, "") @property - def search_query(self): + def search_query(self) -> str: """Returns the Search query.""" - return self.get_attribute("search_query") + return str(self.get_attribute("search_query")) @property - def title(self): + def title(self) -> str: """Returns the Sigma rule title.""" - return self.get_attribute("title") + return str(self.get_attribute("title")) @property - def id(self): + def id(self) -> str: """Returns the Sigma rule id.""" - return self.get_attribute("id") + return str(self.get_attribute("id")) @property - def rule_uuid(self): + def rule_uuid(self) -> str: """Returns the rule id.""" - return self.get_attribute("id") + return str(self.get_attribute("id")) @property - def description(self): + def description(self) -> str: """Returns the rule description.""" - return self.get_attribute("description") + return str(self.get_attribute("description")) @property - def level(self): + def level(self) -> str: """Returns the rule confidence level.""" - return self.get_attribute("level") + return str(self.get_attribute("level")) @property - def falsepositives(self): + def falsepositives(self) -> Any: """Returns the rule falsepositives.""" return self.get_attribute("falsepositives") @property - def author(self): + def author(self) -> str: """Returns the rule author.""" - return self.get_attribute("author") + return str(self.get_attribute("author")) @property - def date(self): + def date(self) -> str: """Returns the rule date.""" - return self.get_attribute("date") + return str(self.get_attribute("date")) @property - def modified(self): + def modified(self) -> str: """Returns the rule modified date.""" - return self.get_attribute("modified") + return str(self.get_attribute("modified")) @property - def logsource(self): + def logsource(self) -> Any: """Returns the rule logsource.""" return self.get_attribute("logsource") @property - def detection(self): + def detection(self) -> Any: """Returns the rule detection.""" return self.get_attribute("detection") @property - def references(self): + def references(self) -> List[str]: """Returns the rule references.""" return self.get_attribute("references") @property - def status(self): + def status(self) -> str: """Returns the rule status.""" - return self.get_attribute("status") + return str(self.get_attribute("status")) - def set_value(self, key, value): + def set_value(self, key: str, value: Any) -> None: """Sets the value for a given key Args: key: key to set the value value: value to set - """ self._attr_dict[key] = value - def _load_rule_dict(self, rule_dict): - """Load a dict into a rule""" + def _load_rule_dict(self, rule_dict: Dict[str, Any]) -> None: + """Load a dict into a rule. + + Args: + rule_dict: Dictionary with rule data. + """ for key, value in rule_dict.items(): self.set_value(key, value) - def from_rule_uuid(self, rule_uuid): + def from_rule_uuid(self, rule_uuid: str) -> None: """Get a SigmaRule object from a rule UUID. Args: rule_uuid: Id of the sigma rule. - """ self.resource_uri = f"sigmarules/{rule_uuid}" @@ -160,7 +169,7 @@ def from_rule_uuid(self, rule_uuid): for key, value in rule_dict.items(): self.set_value(key, value) - def from_text(self, rule_text): + def from_text(self, rule_text: str) -> None: """Obtain a parsed Sigma rule by providing text. Args: @@ -184,7 +193,7 @@ def from_text(self, rule_text): for key, value in rule_dict.items(): self.set_value(key, value) - def delete(self): + def delete(self) -> bool: """Deletes the Sigma rule from Timesketch.""" if not self.get_attribute("id"): logger.warning( diff --git a/api_client/python/timesketch_api_client/sigma_test.py b/api_client/python/timesketch_api_client/sigma_test.py index 7b99f90c7d..5d4f401efe 100644 --- a/api_client/python/timesketch_api_client/sigma_test.py +++ b/api_client/python/timesketch_api_client/sigma_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import test_lib from . import client diff --git a/api_client/python/timesketch_api_client/sketch.py b/api_client/python/timesketch_api_client/sketch.py index ad74426d2a..a1e91db023 100644 --- a/api_client/python/timesketch_api_client/sketch.py +++ b/api_client/python/timesketch_api_client/sketch.py @@ -13,14 +13,14 @@ # limitations under the License. """Timesketch API client library.""" -from __future__ import unicode_literals +from __future__ import annotations import copy import os import json import time import logging -from typing import Dict, Generator, List, Optional, Union +from typing import Any, Dict, Generator, List, Optional, Union, TYPE_CHECKING import pandas @@ -40,6 +40,10 @@ from . import timeline from . import scenario as scenario_lib +if TYPE_CHECKING: + from .client import TimesketchApi + + logger = logging.getLogger("timesketch_api.sketch") @@ -47,17 +51,17 @@ class Sketch(resource.BaseResource): """Timesketch sketch object. A sketch in Timesketch is a collection of one or more timelines. It has - access control and its own namespace for things like labels and comments. - Attributes: - id: The ID of the sketch. - api: An instance of TimesketchApi object. + id (int): The ID of the sketch. + api (TimesketchApi): An instance of TimesketchApi object. """ # Add in necessary fields in data ingested via a different mechanism. _NECESSARY_DATA_FIELDS = frozenset(["timestamp", "datetime", "message"]) - def __init__(self, sketch_id, api, sketch_name=None): + def __init__( + self, sketch_id: int, api: TimesketchApi, sketch_name: Optional[str] = None + ) -> None: """Initializes the Sketch object. Args: @@ -65,6 +69,7 @@ def __init__(self, sketch_id, api, sketch_name=None): api: An instance of TimesketchApi object. sketch_name: Name of the sketch (optional). """ + self.id = sketch_id self.api = api self._archived = None @@ -72,7 +77,7 @@ def __init__(self, sketch_id, api, sketch_name=None): super().__init__(api=api, resource_uri=f"sketches/{self.id}/") @property - def acl(self): + def acl(self) -> Dict[str, Any]: """Property that returns back a ACL dict.""" data = self.lazyload_data(refresh_cache=True) objects = data.get("objects") @@ -85,19 +90,20 @@ def acl(self): return json.loads(permission_string) @property - def attributes(self): + def attributes(self) -> Dict[str, Any]: """Property that returns the sketch attributes.""" data = self.lazyload_data(refresh_cache=True) meta = data.get("meta", {}) return meta.get("attributes", {}) @property - def attributes_table(self): + def attributes_table(self) -> pandas.DataFrame: """DEPRECATED: Property that returns the sketch attributes as a data frame. Given the fluid setup of attributes, this is not a good way to - represent the data. Use the attributes property instead.""" + represent the data. Use the attributes property instead. + """ data = self.lazyload_data(refresh_cache=True) meta = data.get("meta", {}) attributes = meta.get("attributes", []) @@ -108,18 +114,22 @@ def attributes_table(self): return data_frame @property - def description(self): + def description(self) -> str: """Property that returns sketch description. Returns: Sketch description as string. """ - sketch = self.lazyload_data() - return sketch["objects"][0]["description"] + sketch_data = self.lazyload_data() + return sketch_data["objects"][0]["description"] @description.setter - def description(self, description_value): - """Change the sketch description to a new value.""" + def description(self, description_value: str) -> None: + """Change the sketch description to a new value. + + Args: + description_value: The new value for the sketch description. + """ if not isinstance(description_value, str): logger.error("Unable to change the name to a non string value") return @@ -136,7 +146,7 @@ def description(self, description_value): _ = self.lazyload_data(refresh_cache=True) @property - def labels(self): + def labels(self) -> List[str]: """Property that returns the sketch labels.""" data = self.lazyload_data(refresh_cache=True) objects = data.get("objects", []) @@ -151,7 +161,7 @@ def labels(self): return [] @property - def last_activity(self): + def last_activity(self) -> str: """Property that returns the last activity. Returns: @@ -162,7 +172,7 @@ def last_activity(self): return meta.get("last_activity", "") @property - def my_acl(self): + def my_acl(self) -> List[str]: """Property that returns back the ACL for the current user.""" data = self.lazyload_data(refresh_cache=True) objects = data.get("objects") @@ -175,20 +185,24 @@ def my_acl(self): return json.loads(permission_string) @property - def name(self): + def name(self) -> str: """Property that returns sketch name. Returns: Sketch name as string. """ if not self._sketch_name: - sketch = self.lazyload_data() - self._sketch_name = sketch["objects"][0]["name"] + sketch_data = self.lazyload_data() + self._sketch_name = sketch_data["objects"][0]["name"] return self._sketch_name @name.setter - def name(self, name_value): - """Change the name of the sketch to a new value.""" + def name(self, name_value: str) -> None: + """Change the name of the sketch to a new value. + + Args: + name_value: The new name of the sketch. + """ if not isinstance(name_value, str): logger.error("Unable to change the name to a non string value") return @@ -206,7 +220,7 @@ def name(self, name_value): _ = self.lazyload_data(refresh_cache=True) @property - def status(self): + def status(self) -> str: """Property that returns sketch status. Returns: @@ -231,14 +245,16 @@ def status(self): return status_list[0].get("status", "Unknown") - def add_attribute_list(self, name, values, ontology="text"): + def add_attribute_list( + self, name: str, values: List[Any], ontology: str = "text" + ) -> Dict[str, Any]: """Adds or modifies attributes to the sketch. Args: - name (str): The name of the attribute. - values (list): A list of values (in their correct type according + name: The name of the attribute. + values: A list of values (in their correct type according to the ontology). - ontology (str): The ontology (matches with + ontology: The ontology (matches with /data/ontology.yaml), which defines how the attribute is interpreted. @@ -271,13 +287,15 @@ def add_attribute_list(self, name, values, ontology="text"): return error.get_response_json(response, logger) - def add_attribute(self, name, value, ontology="text"): + def add_attribute( + self, name: str, value: Any, ontology: str = "text" + ) -> Dict[str, Any]: """Adds or modifies an attribute to the sketch. Args: - name (str): The name of the attribute. - value (str): Value of the attribute, stored as a string. - ontology (str): The ontology (matches with + name: The name of the attribute. + value: Value of the attribute, stored as a string. + ontology: The ontology (matches with /data/ontology.yaml), which defines how the attribute is interpreted. @@ -292,14 +310,14 @@ def add_attribute(self, name, value, ontology="text"): return self.add_attribute_list(name=name, values=[value], ontology=ontology) - def add_sketch_label(self, label): + def add_sketch_label(self, label: str) -> bool: """Add a label to the sketch. Args: - label (str): A string with the label to add to the sketch. + label: A string with the label to add to the sketch. Returns: - bool: A boolean to indicate whether the label was successfully + A boolean to indicate whether the label was successfully added to the sketch. """ if label in self.labels: @@ -320,12 +338,12 @@ def add_sketch_label(self, label): return status - def remove_attribute(self, name, ontology): + def remove_attribute(self, name: str, ontology: str) -> bool: """Remove an attribute from the sketch. Args: - name (str): The name of the attribute. - ontology (str): The ontology (matches with + name: The name of the attribute. + ontology: The ontology (matches with /data/ontology.yaml), which defines how the attribute is interpreted. @@ -355,14 +373,14 @@ def remove_attribute(self, name, ontology): return status - def remove_sketch_label(self, label): + def remove_sketch_label(self, label: str) -> bool: """Remove a label from the sketch. Args: - label (str): A string with the label to remove from the sketch. + label: A string with the label to remove from the sketch. Returns: - bool: A boolean to indicate whether the label was successfully + A boolean to indicate whether the label was successfully removed from the sketch. """ if label not in self.labels: @@ -386,16 +404,22 @@ def remove_sketch_label(self, label): return status - def create_view(self, name, query_string="", query_dsl="", query_filter=None): + def create_view( + self, + name: str, + query_string: str = "", + query_dsl: str = "", + query_filter: Optional[Dict[str, Any]] = None, + ) -> search.Search: """Create a view object. Args: - name (str): the name of the view. - query_string (str): OpenSearch query string. This is optional + name: the name of the view. + query_string: OpenSearch query string. This is optional yet either a query string or a query DSL is required. - query_dsl (str): OpenSearch query DSL as JSON string. This is + query_dsl: OpenSearch query DSL as JSON string. This is optional yet either a query string or a query DSL is required. - query_filter (dict): Filter for the query as a dict. + query_filter: Filter for the query as a dict. Raises: ValueError: if neither query_string nor query_dsl is provided or @@ -426,11 +450,11 @@ def create_view(self, name, query_string="", query_dsl="", query_filter=None): search_obj.save() return search_obj - def create_story(self, title: str): + def create_story(self, title: str) -> story.Story: """Create a story object. Args: - title (str): the title of the story. + title: the title of the story. Raises: RuntimeError: if a story wasn't created for some reason. @@ -459,7 +483,7 @@ def create_story(self, title: str): story_dict = response_json.get("objects", [{}])[0] return story.Story(story_id=story_dict.get("id", 0), sketch=self, api=self.api) - def delete(self, force_delete=False): + def delete(self, force_delete: bool = False) -> bool: """Deletes the sketch from Timesketch. This method allows for either a soft deletion or a hard deletion @@ -469,7 +493,7 @@ def delete(self, force_delete=False): first be unarchived before a delete operation can be performed. Args: - force_delete (bool): If True, a hard delete is performed, which + force_delete: If True, a hard delete is performed, which permanently removes the sketch and all its associated data (timelines, events, views, etc.) from the Timesketch database and OpenSearch. Administrators can use this to permanently @@ -479,7 +503,7 @@ def delete(self, force_delete=False): indices to free up cluster resources. Returns: - bool: True if the sketch was successfully deleted (either soft or hard). + True if the sketch was successfully deleted (either soft or hard). Raises: RuntimeError: @@ -513,21 +537,21 @@ def delete(self, force_delete=False): def add_to_acl( self, - user_list=None, - group_list=None, - make_public=False, - permissions=None, - ): + user_list: Optional[List[str]] = None, + group_list: Optional[List[str]] = None, + make_public: bool = False, + permissions: Optional[List[str]] = None, + ) -> bool: """Add users or groups to the sketch ACL. Args: - user_list (list[str]): optional list of users to add to the ACL + user_list: optional list of users to add to the ACL of the sketch. Each user is a string. - group_list (list[str]): optional list of groups to add to the ACL + group_list: optional list of groups to add to the ACL of the sketch. Each user is a string. - make_public (bool): Optional boolean indicating the sketch should be + make_public: Optional boolean indicating the sketch should be marked as public. - permissions (list[str]): optional list of permissions (read, write, delete). + permissions: optional list of permissions (read, write, delete). If not the default set of permissions are applied (read, write) Returns: @@ -540,7 +564,7 @@ def add_to_acl( self.api.api_root, self.id ) - data = {} + data: Dict[str, Any] = {} if group_list: group_list_corrected = [str(x).strip() for x in group_list] data["groups"] = group_list_corrected @@ -578,7 +602,7 @@ def add_to_acl( _ = self.lazyload_data(refresh_cache=True) return error.check_return_status(response, logger) - def list_aggregation_groups(self): + def list_aggregation_groups(self) -> List[aggregation.AggregationGroup]: """List all saved aggregation groups for this sketch. Returns: @@ -599,13 +623,17 @@ def list_aggregation_groups(self): groups.append(group) return groups - def list_aggregations(self, include_labels=None, exclude_labels=None): + def list_aggregations( + self, + include_labels: Optional[List[str]] = None, + exclude_labels: Optional[List[str]] = None, + ) -> List[aggregation.Aggregation]: """List all saved aggregations for this sketch. Args: - include_labels (list): list of strings with labels. If defined + include_labels: list of strings with labels. If defined then only return aggregations that have the label in the list. - exclude_labels (list): list of strings with labels. If defined + exclude_labels: list of strings with labels. If defined then only return aggregations that don't have a label in the list. include_labels will be processed first in case both are defined. @@ -653,7 +681,7 @@ def list_aggregations(self, include_labels=None, exclude_labels=None): aggregations.append(aggregation_obj) return aggregations - def list_graphs(self): + def list_graphs(self) -> List[graph.Graph]: """Returns a list of stored graphs.""" if self.is_archived(): raise RuntimeError("Unable to list graphs on an archived sketch.") @@ -675,11 +703,13 @@ def list_graphs(self): return_list.append(graph_obj) return return_list - def get_analyzer_status(self, as_sessions=False): + def get_analyzer_status( + self, as_sessions: bool = False + ) -> Union[List[analyzer.AnalyzerResult], List[Dict[str, Any]]]: """Returns a list of started analyzers and their status. Args: - as_sessions (bool): optional, if set to True then a list of + as_sessions: optional, if set to True then a list of AnalyzerResult objects will be returned. Defaults to returning a list of dicts. Returns: @@ -723,7 +753,7 @@ def get_analyzer_status(self, as_sessions=False): return stats_list - def get_aggregation(self, aggregation_id): + def get_aggregation(self, aggregation_id: int) -> Optional[aggregation.Aggregation]: """Return a stored aggregation. Args: @@ -740,7 +770,9 @@ def get_aggregation(self, aggregation_id): return aggregation_obj return None - def get_aggregation_group(self, group_id): + def get_aggregation_group( + self, group_id: int + ) -> Optional[aggregation.AggregationGroup]: """Return a stored aggregation group. Args: @@ -760,13 +792,15 @@ def get_aggregation_group(self, group_id): return group_obj return None - def get_story(self, story_id=None, story_title=None): + def get_story( + self, story_id: Optional[int] = None, story_title: Optional[str] = None + ) -> Optional[story.Story]: """Returns a story object that is stored in the sketch. Args: - story_id (int): an integer indicating the ID of the story to + story_id: an integer indicating the ID of the story to be fetched. Defaults to None. - story_title (str): a string with the title of the story. Optional + story_title: a string with the title of the story. Optional and defaults to None. Returns: @@ -789,7 +823,9 @@ def get_story(self, story_id=None, story_title=None): return story_obj return None - def get_view(self, view_id=None, view_name=None): + def get_view( + self, view_id: Optional[int] = None, view_name: Optional[str] = None + ) -> Optional[search.Search]: """Returns a saved search object that is stored in the sketch. Args: @@ -810,13 +846,15 @@ def get_view(self, view_id=None, view_name=None): return self.get_saved_search(search_id=view_id, search_name=view_name) - def get_saved_search(self, search_id=None, search_name=None): + def get_saved_search( + self, search_id: Optional[int] = None, search_name: Optional[str] = None + ) -> Optional[search.Search]: """Returns a saved search object that is stored in the sketch. Args: - search_id (int): an integer indicating the ID of the saved search to + search_id: an integer indicating the ID of the saved search to be fetched. Defaults to None. - search_name (str): a string with the name of the saved search. Optional + search_name: a string with the name of the saved search. Optional and defaults to None. Returns: @@ -837,13 +875,15 @@ def get_saved_search(self, search_id=None, search_name=None): return search_obj return None - def get_timeline(self, timeline_id=None, timeline_name=None): + def get_timeline( + self, timeline_id: Optional[int] = None, timeline_name: Optional[str] = None + ) -> Optional[timeline.Timeline]: """Returns a timeline object that is stored in the sketch. Args: - timeline_id (int): an integer indicating the ID of the timeline to + timeline_id: an integer indicating the ID of the timeline to be fetched. Defaults to None. - timeline_name (str): a string with the name of the timeline. Optional + timeline_name: a string with the name of the timeline. Optional and defaults to None. Returns: @@ -865,7 +905,7 @@ def get_timeline(self, timeline_id=None, timeline_name=None): return timeline_ return None - def get_intelligence_attribute(self): + def get_intelligence_attribute(self) -> Dict[str, Any]: """Returns a timeline object that is stored in the sketch. Returns: @@ -881,7 +921,7 @@ def get_intelligence_attribute(self): return intel_attribute - def list_stories(self): + def list_stories(self) -> List[story.Story]: """Get a list of all stories that are attached to the sketch. Returns: @@ -913,7 +953,7 @@ def list_stories(self): ) return story_list - def list_views(self): + def list_views(self) -> List[search.Search]: """List all saved views for this sketch. Returns: @@ -925,7 +965,7 @@ def list_views(self): ) return self.list_saved_searches() - def list_saved_searches(self): + def list_saved_searches(self) -> List[search.Search]: """List all saved searches for this sketch. Returns: @@ -952,7 +992,7 @@ def list_saved_searches(self): return searches - def list_search_templates(self): + def list_search_templates(self) -> List[searchtemplate.SearchTemplate]: """Get a list of all search templates that are available. Returns: @@ -974,7 +1014,7 @@ def list_search_templates(self): return template_list - def list_timelines(self): + def list_timelines(self) -> List[timeline.Timeline]: """List all timelines for this sketch. Returns: @@ -1002,7 +1042,7 @@ def list_timelines(self): return timelines # pylint: disable=unused-argument - def add_timeline(self, searchindex): + def add_timeline(self, searchindex: Any) -> None: """Deprecated function to add timeline to sketch. Args: @@ -1022,42 +1062,42 @@ def add_timeline(self, searchindex): # pylint: disable=too-many-arguments def explore( self, - query_string=None, - query_dsl=None, - query_filter=None, - view=None, - return_fields=None, - as_pandas=False, - max_entries=None, - file_name="", - as_object=False, + query_string: Optional[str] = None, + query_dsl: Optional[str] = None, + query_filter: Optional[Dict[str, Any]] = None, + view: Optional[search.Search] = None, + return_fields: Optional[str] = None, + as_pandas: bool = False, + max_entries: Optional[int] = None, + file_name: str = "", + as_object: bool = False, use_wildcard_fields: bool = False, - ): + ) -> Union[Dict[str, Any], pandas.DataFrame, search.Search, None]: """Explore the sketch. Args: - query_string (str): OpenSearch query string. - query_dsl (str): OpenSearch query DSL as JSON string. - query_filter (dict): Filter for the query as a dict. - view (search.Search): View object instance (optional). - return_fields (str): A comma separated string with a list of fields + query_string: OpenSearch query string. + query_dsl: OpenSearch query DSL as JSON string. + query_filter: Filter for the query as a dict. + view: View object instance (optional). + return_fields: A comma separated string with a list of fields that should be included in the response. Optional and defaults to None. - as_pandas (bool): Optional bool that determines if the results + as_pandas: Optional bool that determines if the results should be returned back as a dictionary or a Pandas DataFrame. - max_entries (int): Optional integer denoting a best effort to limit + max_entries: Optional integer denoting a best effort to limit the output size to the number of events. Events are read in, 10k at a time so there may be more events in the answer back than this number denotes, this is a best effort. - file_name (str): Optional filename, if provided the results of + file_name: Optional filename, if provided the results of the query will be exported to a ZIP file instead of being returned back as a dict or a pandas DataFrame. The ZIP file will contain a METADATA file and a CSV with the results from the query. - as_object (bool): Optional bool that determines whether the + as_object: Optional bool that determines whether the function will return a search object back instead of raw results. - use_wildcard_fields (bool): Optional bool, if set to True compiles + use_wildcard_fields: Optional bool, if set to True compiles the search query using native wildcard fields mapping. Returns: @@ -1103,7 +1143,8 @@ def explore( return search_obj if file_name: - return search_obj.to_file(file_name) + search_obj.to_file(file_name) + return None if as_pandas: return search_obj.to_pandas() @@ -1114,7 +1155,7 @@ def explore_wildcard( self, query_string: str, limit: Optional[int] = None, - ) -> Dict[str, Union[Dict, List]]: + ) -> Dict[str, Union[Dict[str, Any], List[Any]]]: """Explore the sketch with raw wildcard queries (Deprecated). This method is maintained for backward-compatibility. Newer scripts should @@ -1133,7 +1174,7 @@ def explore_wildcard( raise RuntimeError("Unable to query an archived sketch.") resource_url = f"{self.api.api_root}/sketches/{self.id}/explore_wildcard/" - form_data = { + form_data: Dict[str, Any] = { "query": query_string, } if limit is not None: @@ -1141,7 +1182,7 @@ def explore_wildcard( response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def list_available_analyzers(self): + def list_available_analyzers(self) -> List[str]: """Returns a list of available analyzers.""" resource_url = "{0:s}/sketches/{1:d}/analyzer/".format( self.api.api_root, self.id @@ -1153,22 +1194,22 @@ def list_available_analyzers(self): def run_analyzer( self, - analyzer_name, - analyzer_kwargs=None, - timeline_id=None, - timeline_name=None, - ): + analyzer_name: str, + analyzer_kwargs: Optional[Dict[str, Any]] = None, + timeline_id: Optional[int] = None, + timeline_name: Optional[str] = None, + ) -> Union[analyzer.AnalyzerResult, str]: """Run an analyzer on a timeline. Args: - analyzer_name (str): the name of the analyzer class to run against the + analyzer_name: the name of the analyzer class to run against the timeline. - analyzer_kwargs (dict): optional dict with parameters for the analyzer. + analyzer_kwargs: optional dict with parameters for the analyzer. This is optional and just for those analyzers that can accept further parameters. - timeline_id (int): the ID of the timeline. This is optional and only + timeline_id: the ID of the timeline. This is optional and only required if timeline_name is not set. - timeline_name (str): the name of the timeline in the timesketch UI. This + timeline_name: the name of the timeline in the timesketch UI. This is optional and only required if timeline_id is not set. If there are more than a single timeline with the same name a timeline_id is required. @@ -1197,9 +1238,9 @@ def run_analyzer( return "Unable to run analyzer, need to define either timeline ID or name" if timeline_name: - sketch = self.lazyload_data(refresh_cache=True) + sketch_data = self.lazyload_data(refresh_cache=True) timelines = [] - for timeline_dict in sketch["objects"][0]["timelines"]: + for timeline_dict in sketch_data["objects"][0]["timelines"]: name = timeline_dict.get("name", "") if timeline_name.lower() == name.lower(): timelines.append(timeline_dict.get("id")) @@ -1233,21 +1274,21 @@ def run_analyzer( def remove_acl( self, - user_list=None, - group_list=None, - remove_public=False, - permissions=None, - ): + user_list: Optional[List[str]] = None, + group_list: Optional[List[str]] = None, + remove_public: bool = False, + permissions: Optional[List[str]] = None, + ) -> bool: """Remove users or groups to the sketch ACL. Args: - user_list (list[str]): optional list of users to remove from the ACL + user_list: optional list of users to remove from the ACL of the sketch. Each user is a string. - group_list (list[str]): optional list of groups to remove from the ACL + group_list: optional list of groups to remove from the ACL of the sketch. Each user is a string. - remove_public (bool): Optional boolean indicating the sketch should be + remove_public: Optional boolean indicating the sketch should be no longer marked as public. - permissions (list[str]): optional list of permissions (read, write, delete). + permissions: optional list of permissions (read, write, delete). If not the default set of permissions are applied (read, write) Returns: @@ -1260,7 +1301,7 @@ def remove_acl( self.api.api_root, self.id ) - data = {} + data: Dict[str, Any] = {} if group_list: group_list_corrected = [str(x).strip() for x in group_list] data["remove_groups"] = group_list_corrected @@ -1274,8 +1315,8 @@ def remove_acl( if permissions: allowed_permissions = set(["read", "write", "delete"]) - permissions = list(allowed_permissions.intersection(set(permissions))) - data["permissions"] = json.dumps(permissions) + permissions_list = list(allowed_permissions.intersection(set(permissions))) + data["permissions"] = json.dumps(permissions_list) if not data: return True @@ -1285,11 +1326,11 @@ def remove_acl( _ = self.lazyload_data(refresh_cache=True) return error.check_return_status(response, logger) - def aggregate(self, aggregate_dsl): + def aggregate(self, aggregate_dsl: str) -> aggregation.Aggregation: """Run an aggregation request on the sketch. Args: - aggregate_dsl (str): OpenSearch aggregation query DSL string. + aggregate_dsl: OpenSearch aggregation query DSL string. Returns: An aggregation object (instance of Aggregation). @@ -1309,7 +1350,7 @@ def aggregate(self, aggregate_dsl): return aggregation_obj - def list_available_aggregators(self): + def list_available_aggregators(self) -> pandas.DataFrame: """Return a list of all available aggregators in the sketch.""" data = self.lazyload_data() meta = data.get("meta", {}) @@ -1345,7 +1386,9 @@ def list_available_aggregators(self): return pandas.DataFrame(entries) - def run_aggregator(self, aggregator_name, aggregator_parameters): + def run_aggregator( + self, aggregator_name: str, aggregator_parameters: Dict[str, Any] + ) -> aggregation.Aggregation: """Run an aggregator class. Args: @@ -1369,20 +1412,20 @@ def run_aggregator(self, aggregator_name, aggregator_parameters): def store_aggregation( self, - name, - description, - aggregator_name, - aggregator_parameters, - chart_type="", - ): + name: str, + description: str, + aggregator_name: str, + aggregator_parameters: Dict[str, Any], + chart_type: str = "", + ) -> Optional[aggregation.Aggregation]: """Store an aggregation in the sketch. Args: - name (str): a name that will be associated with the aggregation. - description (str): description of the aggregation, visible in the UI. - aggregator_name (str): name of the aggregator class. - aggregator_parameters (dict): parameters of the aggregator. - chart_type (str): string representing the chart type. + name: a name that will be associated with the aggregation. + description: description of the aggregation, visible in the UI. + aggregator_name: name of the aggregator class. + aggregator_parameters: parameters of the aggregator. + chart_type: string representing the chart type. Raises: RuntimeError: if the client is unable to store the aggregation. @@ -1410,13 +1453,15 @@ def store_aggregation( return None - def comment_event(self, event_id, index, comment_text): + def comment_event( + self, event_id: str, index: str, comment_text: str + ) -> Dict[str, Any]: """Adds a comment to a single event. Args: - event_id (str): id of the event - index (str): The OpenSearch index name - comment_text (str): text to add as a comment + event_id: id of the event + index: The OpenSearch index name + comment_text: text to add as a comment Returns: a json data of the query. """ @@ -1438,7 +1483,7 @@ def comment_event(self, event_id, index, comment_text): response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def add_event_attributes(self, events): + def add_event_attributes(self, events: List[Dict[str, Any]]) -> Dict[str, Any]: """Add attributes to one or more events. Args: @@ -1461,7 +1506,7 @@ def add_event_attributes(self, events): return error.get_response_json(response, logger) - def get_event(self, event_id, index_id): + def get_event(self, event_id: Union[int, str], index_id: str) -> Dict[str, Any]: """Gets information about an event, including raw event and meta data. Args: @@ -1479,20 +1524,26 @@ def get_event(self, event_id, index_id): ) resource_url_params = "?searchindex_id={0:s}&event_id={1:s}".format( - index_id, event_id + str(index_id), str(event_id) ) response = self.api.session.get(resource_url_base + resource_url_params) return error.get_response_json(response, logger) - def label_events(self, events, label_name, remove=False, conclusion_id=None): + def label_events( + self, + events: List[Dict[str, Any]], + label_name: str, + remove: bool = False, + conclusion_id: Optional[int] = None, + ) -> Dict[str, Any]: """Labels one or more events with label_name. Args: - events (json): Array of JSON objects representing events. - label_name (string): String to label the event with. - remove (bool): If true, the label will be removed instead of added. - conclusion_id (int): Optional. ID of a conclusion to link the + events: Array of JSON objects representing events. + label_name: String to label the event with. + remove: If true, the label will be removed instead of added. + conclusion_id: Optional. ID of a conclusion to link the event to. Returns: @@ -1501,7 +1552,7 @@ def label_events(self, events, label_name, remove=False, conclusion_id=None): if self.is_archived(): raise RuntimeError("Unable to label events in an archived sketch.") - form_data = { + form_data: Dict[str, Any] = { "annotation": label_name, "annotation_type": "label", "events": events, @@ -1516,13 +1567,15 @@ def label_events(self, events, label_name, remove=False, conclusion_id=None): response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def link_event_to_conclusion(self, events, conclusion_id, unlink=False): + def link_event_to_conclusion( + self, events: List[Dict[str, Any]], conclusion_id: int, unlink: bool = False + ) -> Dict[str, Any]: """Links one or more events to a conclusion as a fact. Args: events: Array of JSON objects representing events. - conclusion_id (int): ID of the conclusion to link the event to. - unlink (bool): If true, the link will be removed. + conclusion_id: ID of the conclusion to link the event to. + unlink: If true, the link will be removed. Returns: Dictionary with query results. @@ -1534,13 +1587,15 @@ def link_event_to_conclusion(self, events, conclusion_id, unlink=False): conclusion_id=conclusion_id, ) - def untag_events(self, events, tags_to_remove: list): + def untag_events( + self, events: List[Dict[str, Any]], tags_to_remove: List[str] + ) -> Dict[str, Any]: """Removes a list of tags from a list of events. The upper limit is 500 (events or tags) based on the API. Args: - events (list): events dict. Must have the structure: + events: events dict. Must have the structure: "events": [ { "_id": event_id, @@ -1564,7 +1619,7 @@ def untag_events(self, events, tags_to_remove: list): response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def untag_event(self, event_id: str, index, tag: str): + def untag_event(self, event_id: str, index: str, tag: str) -> Dict[str, Any]: """Removes a tag from an event. This method can be used if just one events needs to be untagged. @@ -1576,7 +1631,7 @@ def untag_event(self, event_id: str, index, tag: str): Args: event_id: id of the event - index (str): The OpenSearch index name + index: The OpenSearch index name tag: tag to remove Returns: @@ -1600,13 +1655,15 @@ def untag_event(self, event_id: str, index, tag: str): response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def tag_events(self, events, tags, verbose=False): + def tag_events( + self, events: List[Dict[str, Any]], tags: List[str], verbose: bool = False + ) -> Dict[str, Any]: """Tags one or more events with a list of tags. Args: - events (list): Array of JSON objects representing events. - tags (list[str]): List of tags (str) to add to the events. - verbose (bool): Bool that determines whether extra information + events: Array of JSON objects representing events. + tags: List of tags (str) to add to the events. + verbose: Bool that determines whether extra information is added to the meta dict that gets returned. Raises: @@ -1648,20 +1705,24 @@ def tag_events(self, events, tags, verbose=False): return meta def search_by_label( - self, label_name, return_fields=None, max_entries=None, as_pandas=False - ): + self, + label_name: str, + return_fields: Optional[str] = None, + max_entries: Optional[int] = None, + as_pandas: bool = False, + ) -> Union[Dict[str, Any], pandas.DataFrame, search.Search, None]: """Searches for all events containing a given label. Args: - label_name (str): A string representing the label to search for. - return_fields (str): A comma separated string with a list of fields + label_name: A string representing the label to search for. + return_fields: A comma separated string with a list of fields that should be included in the response. Optional and defaults to None. - max_entries (int): Optional integer denoting a best effort to limit + max_entries: Optional integer denoting a best effort to limit the output size to the number of events. Events are read in, 10k at a time so there may be more events in the answer back than this number denotes, this is a best effort. - as_pandas (bool): Optional bool that determines if the results should + as_pandas: Optional bool that determines if the results should be returned back as a dictionary or a Pandas DataFrame. Returns: @@ -1695,13 +1756,18 @@ def search_by_label( as_pandas=as_pandas, ) - def add_scenario(self, uuid=None, dfiq_id=None, name=None): + def add_scenario( + self, + uuid: Optional[str] = None, + dfiq_id: Optional[str] = None, + name: Optional[str] = None, + ) -> scenario_lib.Scenario: """Adds an investigative scenario to the sketch. Args: - uuid (str): [Optional] UUID of the DFIQ scenario template to add. - dfiq_id (str): [Optional] ID of the DFIQ scenario template to add. - name (str): [Optional] Name of the scenario to add. + uuid: [Optional] UUID of the DFIQ scenario template to add. + dfiq_id: [Optional] ID of the DFIQ scenario template to add. + name: [Optional] Name of the scenario to add. Raises: ValueError: If none or more than one of uuid, dfiq_id, or name are provided. @@ -1721,7 +1787,7 @@ def add_scenario(self, uuid=None, dfiq_id=None, name=None): "Exactly one of 'uuid', 'dfiq_id', or 'name' must be provided." ) - form_data = {} + form_data: Dict[str, Any] = {} if uuid: form_data["uuid"] = uuid elif dfiq_id: @@ -1754,7 +1820,7 @@ def add_scenario(self, uuid=None, dfiq_id=None, name=None): api=self.api, ) - def list_scenarios(self): + def list_scenarios(self) -> List[scenario_lib.Scenario]: """Get a list of all scenarios that are attached to the sketch. Returns: @@ -1782,13 +1848,18 @@ def list_scenarios(self): for scenario_data in scenario_objects ] - def add_question(self, dfiq_id=None, uuid=None, question_text=None): + def add_question( + self, + dfiq_id: Optional[str] = None, + uuid: Optional[str] = None, + question_text: Optional[str] = None, + ) -> scenario_lib.Question: """Adds an investigative question to the sketch. Args: - dfiq_id (str): [Optional] ID of the DFIQ question template to add. - uuid (str): [Optional] UUID of the DFIQ question template to add. - question_text (str): [Optional] Question text to add. + dfiq_id: [Optional] ID of the DFIQ question template to add. + uuid: [Optional] UUID of the DFIQ question template to add. + question_text: [Optional] Question text to add. Raises: ValueError: If none or more than one of dfiq_id, uuid, or @@ -1809,7 +1880,7 @@ def add_question(self, dfiq_id=None, uuid=None, question_text=None): "provided." ) - form_data = {} + form_data: Dict[str, Any] = {} if dfiq_id: form_data["template_id"] = dfiq_id elif uuid: @@ -1842,7 +1913,7 @@ def add_question(self, dfiq_id=None, uuid=None, question_text=None): api=self.api, ) - def list_questions(self): + def list_questions(self) -> List[scenario_lib.Question]: """Get a list of all questions attached to the sketch. Returns: @@ -1871,16 +1942,23 @@ def list_questions(self): for question_data in question_objects ] - def add_event(self, message, date, timestamp_desc, attributes=None, tags=None): + def add_event( + self, + message: str, + date: str, + timestamp_desc: str, + attributes: Optional[Dict[str, Any]] = None, + tags: Optional[List[str]] = None, + ) -> Dict[str, Any]: """Adds an event to the sketch specific timeline. Args: - message (str): A string that will be used as the message string. - date (str): A string with the timestamp of the message. This should be + message: A string that will be used as the message string. + date: A string with the timestamp of the message. This should be in a human readable format, eg: "2020-09-03T22:52:21". - timestamp_desc (str): Description of the timestamp. - attributes (dict): A dict of extra attributes to add to the event. - tags (list[str]): A list of strings to include as tags. + timestamp_desc: Description of the timestamp. + attributes: A dict of extra attributes to add to the event. + tags: A list of strings to include as tags. Raises: ValueError: If tags is not a list of strings or attributes @@ -1908,7 +1986,7 @@ def add_event(self, message, date, timestamp_desc, attributes=None, tags=None): if not isinstance(attributes, dict): raise ValueError("Attributes needs to be a dict.") - form_data = { + form_data: Dict[str, Any] = { "date_string": date, "timestamp_desc": timestamp_desc, "message": message, @@ -1932,7 +2010,7 @@ def add_event(self, message, date, timestamp_desc, attributes=None, tags=None): response = self.api.session.post(resource_url, json=form_data) return error.get_response_json(response, logger) - def is_archived(self): + def is_archived(self) -> bool: """Return a boolean indicating whether the sketch has been archived.""" if self._archived is not None: return self._archived @@ -1946,7 +2024,7 @@ def is_archived(self): self._archived = meta.get("is_archived", False) return self._archived - def archive(self): + def archive(self) -> bool: """Archive a sketch and return a boolean whether it was successful.""" if self.is_archived(): logger.error("Sketch already archived.") @@ -1962,7 +2040,7 @@ def archive(self): return return_status - def unarchive(self): + def unarchive(self) -> bool: """Unarchives a sketch and return boolean whether it was successful.""" if not self.is_archived(): logger.error("Sketch wasn't archived.") @@ -1980,12 +2058,12 @@ def unarchive(self): self._archived = not return_status return return_status - def export(self, file_path, stream=False): + def export(self, file_path: str, stream: bool = False) -> None: """Exports the content of the sketch to a ZIP file. Args: - file_path (str): a file path where the ZIP file will be saved. - stream (bool): whether to stream the download. + file_path: a file path where the ZIP file will be saved. + stream: whether to stream the download. Raises: RuntimeError: if sketch cannot be exported. @@ -2039,21 +2117,21 @@ def export_events_stream( self, query_string: Optional[str] = None, query_dsl: Optional[str] = None, - query_filter: Optional[Dict] = None, + query_filter: Optional[Dict[str, Any]] = None, return_fields: Optional[List[str]] = None, - ) -> Generator[Dict, None, None]: + ) -> Generator[Dict[str, Any], None, None]: """Exports all events from the sketch matching the query. This uses the high-performance sliced export API endpoint. Args: - query_string (str): OpenSearch query string. - query_dsl (str): OpenSearch query DSL as JSON string. - query_filter (dict): Filter for the query as a dict. - return_fields (list): List of strings with fields to return. + query_string: OpenSearch query string. + query_dsl: OpenSearch query DSL as JSON string. + query_filter: Filter for the query as a dict. + return_fields: List of strings with fields to return. Yields: - dict: A dictionary representing an event. + A dictionary representing an event. """ if return_fields is None: return_fields = ["datetime", "message", "timestamp_desc"] @@ -2063,14 +2141,18 @@ def export_events_stream( if not (query_string or query_filter or query_dsl): query_string = "*" - if return_fields and isinstance(return_fields, list): - return_fields = ",".join(return_fields) + fields_string = None + if return_fields: + if isinstance(return_fields, list): + fields_string = ",".join(return_fields) + else: + fields_string = return_fields form_data = { "query": query_string, "filter": query_filter, "dsl": query_dsl, - "fields": return_fields, + "fields": fields_string, } response = self.api.session.post(resource_url, json=form_data, stream=True) @@ -2088,7 +2170,9 @@ def export_events_stream( logger.warning("Received invalid JSON line during export") continue - def create_timeline(self, searchindex_id: int, timeline_name: str): + def create_timeline( + self, searchindex_id: int, timeline_name: str + ) -> timeline.Timeline: """Creates a Timeline in this Sketch This method attempts to create a new timeline associated with this sketch @@ -2098,9 +2182,9 @@ def create_timeline(self, searchindex_id: int, timeline_name: str): API errors, or unexpected response formats. Args: - searchindex_id (int): The ID of the SearchIndex that holds the data + searchindex_id: The ID of the SearchIndex that holds the data for this timeline. - timeline_name (str): The name of the timeline + timeline_name: The name of the timeline Returns: An instance of a Timeline object representing the newly created @@ -2117,7 +2201,7 @@ def create_timeline(self, searchindex_id: int, timeline_name: str): resource_url = f"{self.api.api_root}/sketches/{self.id}/timelines/" form_data = {"timeline": searchindex_id, "timeline_name": timeline_name} - last_exception = None + last_exception: Optional[Exception] = None for attempt in range(self.api.DEFAULT_RETRY_COUNT): try: @@ -2210,14 +2294,14 @@ def create_timeline(self, searchindex_id: int, timeline_name: str): def create_datasource( self, timeline_id: int, provider: str, context: str, data_label: str - ): + ) -> Dict[str, Any]: """Creates a datasource Args: - timeline_id (int): id of the Timeline that this datasource is part of. - provider (str): Name of the application that collected the data. - context (str): Context on how the data was collected. - data_label (str): Data label for the uploaded data. + timeline_id: id of the Timeline that this datasource is part of. + provider: Name of the application that collected the data. + context: Context on how the data was collected. + data_label: Data label for the uploaded data. Raises: ValueError: If the datasource object fails to create @@ -2245,15 +2329,15 @@ def create_datasource( def generate_timeline_from_es_index( self, - es_index_name, - name, - index_name="", - description="", - provider="Manually added to OpenSearch", - context="Added via API client", - data_label="OpenSearch", - status="ready", - ): + es_index_name: str, + name: str, + index_name: str = "", + description: str = "", + provider: str = "Manually added to OpenSearch", + context: str = "Added via API client", + data_label: str = "OpenSearch", + status: str = "ready", + ) -> timeline.Timeline: """Creates and returns a Timeline from OpenSearch data. This function can be used to import data into a sketch that was @@ -2263,19 +2347,19 @@ def generate_timeline_from_es_index( Timeline) for Timesketch to be able to properly support it. Args: - es_index_name (str): name of the index in OpenSearch. - name (str): string with the name of the timeline. - index_name (str): optional string for the SearchIndex name, defaults + es_index_name: name of the index in OpenSearch. + name: string with the name of the timeline. + index_name: optional string for the SearchIndex name, defaults to the same as the es_index_name. - description (str): optional string with a description of the timeline. - provider (str): optional string with the provider name for the data + description: optional string with a description of the timeline. + provider: optional string with the provider name for the data source of the imported data. Defaults to "Manually added to OpenSearch". - context (str): optional string with the context for the data upload, + context: optional string with the context for the data upload, defaults to "Added via API client". - data_label (str): optional string with the data label of the OpenSearch + data_label: optional string with the data label of the OpenSearch data, defaults to "OpenSearch". - status (str): Optional string, if provided will be used as a status + status: Optional string, if provided will be used as a status for the searchindex, valid options are: "ready", "fail", "processing", "timeout". Defaults to "ready". @@ -2344,15 +2428,21 @@ def generate_timeline_from_es_index( return created_timeline - def run_data_finder(self, start_date, end_date, rule_names, timelines=None): + def run_data_finder( + self, + start_date: str, + end_date: str, + rule_names: List[str], + timelines: Optional[List[Union[int, str]]] = None, + ) -> List[Dict[str, Any]]: """Runs the data finder . Args: - start_date (str): Start date as a ISO 8601 formatted string. - end_date (str): End date as a ISO 8601 formatted string. - rule_names (list): A list of strings with rule names to run + start_date: Start date as a ISO 8601 formatted string. + end_date: End date as a ISO 8601 formatted string. + rule_names: A list of strings with rule names to run against the dataset in the sketch. - timelines (list): Optional list of timeline identifiers or + timelines: Optional list of timeline identifiers or timeline names to limit the data search to certain timelines within the sketch. Defaults to search all timelines. diff --git a/api_client/python/timesketch_api_client/sketch_test.py b/api_client/python/timesketch_api_client/sketch_test.py index 7921cd2779..b101d147c8 100644 --- a/api_client/python/timesketch_api_client/sketch_test.py +++ b/api_client/python/timesketch_api_client/sketch_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import client from . import search diff --git a/api_client/python/timesketch_api_client/story.py b/api_client/python/timesketch_api_client/story.py index 8d9d40fa7a..f7e04dceed 100644 --- a/api_client/python/timesketch_api_client/story.py +++ b/api_client/python/timesketch_api_client/story.py @@ -13,8 +13,11 @@ # limitations under the License. """Timesketch API client library.""" +from __future__ import annotations + import json import logging +from typing import Any, Dict, List, Optional, TYPE_CHECKING, Union import pandas as pd @@ -23,6 +26,11 @@ from . import resource from . import search +if TYPE_CHECKING: + from .client import TimesketchApi + from .sketch import Sketch + from .view import View + logger = logging.getLogger("timesketch_api.story") @@ -32,30 +40,39 @@ class BaseBlock: # A string representation of the type of block. TYPE = "" - def __init__(self, story, index): - """Initialize the base block.""" + def __init__(self, story: Story, index: int) -> None: + """Initialize the base block. + + Args: + story: Story object. + index: Index of the block. + """ self.index = index self._story = story - self._data = None + self._data: Any = None @property - def data(self): + def data(self) -> Union[Dict[str, Any], str, List[Any], None]: """Returns the building block.""" return self.to_dict() @data.setter - def data(self, data): - """Feeds data to the block.""" + def data(self, data: Any) -> None: + """Feeds data to the block. + + Args: + data: Data to feed to the block. + """ self._data = data self._story.commit() @property - def json(self): + def json(self) -> Union[Dict[str, Any], str, List[Any], None]: """Returns the building block.""" return self.to_dict() - def _get_base(self): + def _get_base(self) -> Dict[str, Any]: """Returns a base building block.""" return { "componentName": "", @@ -66,31 +83,39 @@ def _get_base(self): "isActive": False, } - def delete(self): + def delete(self) -> None: """Remove block from index.""" self._story.remove_block(self.index) - def feed(self, data): - """Feed data into the block.""" + def feed(self, data: Any) -> None: + """Feed data into the block. + + Args: + data: Data to feed into the block. + """ self._data = data - def from_dict(self, data_dict): - """Feed a block from a block dict.""" + def from_dict(self, data_dict: Dict[str, Any]) -> None: + """Feed a block from a block dict. + + Args: + data_dict: Dictionary with block data. + """ raise NotImplementedError - def move_down(self): + def move_down(self) -> None: """Moves a block down one location in the index.""" new_index = self.index + 1 self._story.move_to(self, new_index) self.index = new_index - def move_up(self): + def move_up(self) -> None: """Moves a block up one location in the index.""" new_index = self.index - 1 self._story.move_to(self, new_index) self.index = new_index - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the block data. Raises: @@ -102,7 +127,7 @@ def to_dict(self): """ raise NotImplementedError - def reset(self): + def reset(self) -> None: """Resets the data in the block.""" self._data = None @@ -112,37 +137,51 @@ class ViewBlock(BaseBlock): TYPE = "view" - def __init__(self, story, index): + def __init__(self, story: Story, index: int) -> None: + """Initialize the view block. + + Args: + story: Story object. + index: Index of the block. + """ super().__init__(story, index) self._view_id = 0 self._view_name = "" @property - def view(self): + def view(self) -> Any: """Returns the view.""" return self._data @view.setter - def view(self, new_view): - """Sets a new view to the block.""" + def view(self, new_view: View) -> None: + """Sets a new view to the block. + + Args: + new_view: View object. + """ self.data = new_view @property - def view_id(self): + def view_id(self) -> int: """Returns the view ID.""" if self._data and hasattr(self._data, "id"): - return self._data.id + return int(self._data.id) return self._view_id @property - def view_name(self): + def view_name(self) -> str: """Returns the view name.""" if self._data and hasattr(self._data, "name"): - return self._data.name + return str(self._data.name) return self._view_name - def from_dict(self, data_dict): - """Feed a block from a block dict.""" + def from_dict(self, data_dict: Dict[str, Any]) -> None: + """Feed a block from a block dict. + + Args: + data_dict: Dictionary with block data. + """ props = data_dict.get("componentProps", {}) view_dict = props.get("view") if not view_dict: @@ -151,7 +190,7 @@ def from_dict(self, data_dict): self._view_id = view_dict.get("id", 0) self._view_name = view_dict.get("name", "") - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the block data. Raises: @@ -187,23 +226,31 @@ class TextBlock(BaseBlock): TYPE = "text" @property - def text(self): + def text(self) -> str: """Returns the text.""" if not self._data: return "" - return self._data + return str(self._data) @text.setter - def text(self, new_text): - """Sets a new text to the block.""" + def text(self, new_text: str) -> None: + """Sets a new text to the block. + + Args: + new_text: New text to set. + """ self.data = new_text - def from_dict(self, data_dict): - """Feed a block from a block dict.""" + def from_dict(self, data_dict: Dict[str, Any]) -> None: + """Feed a block from a block dict. + + Args: + data_dict: Dictionary with block data. + """ text = data_dict.get("content", "") self.feed(text) - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the block data. Raises: @@ -230,61 +277,79 @@ class AggregationBlock(BaseBlock): TYPE = "aggregation" - def __init__(self, story, index): + def __init__(self, story: Story, index: int) -> None: + """Initialize the aggregation block. + + Args: + story: Story object. + index: Index of the block. + """ super().__init__(story, index) self._agg_id = 0 self._agg_name = "" - self._agg_dict = {} + self._agg_dict: Dict[str, Any] = {} self._chart_type = "table" @property - def aggregation(self): + def aggregation(self) -> Optional[aggregation.Aggregation]: """Returns the aggregation object.""" if self._data: return self._data return None @aggregation.setter - def aggregation(self, agg_obj): - """Set the aggregation object.""" + def aggregation(self, agg_obj: aggregation.Aggregation) -> None: + """Set the aggregation object. + + Args: + agg_obj: Aggregation object. + """ self._data = agg_obj @property - def agg_name(self): + def agg_name(self) -> str: """Returns the aggregation name.""" return self._agg_name @property - def agg_id(self): + def agg_id(self) -> int: """Returns the aggregation ID.""" return self._agg_id @property - def chart_type(self): + def chart_type(self) -> str: """Returns the aggregation type.""" return self._chart_type @chart_type.setter - def chart_type(self, new_type): - """Sets the aggregation type.""" + def chart_type(self, new_type: str) -> None: + """Sets the aggregation type. + + Args: + new_type: Type of the chart. + """ self._chart_type = new_type @property - def table(self): + def table(self) -> pd.DataFrame: """Returns a table view, as a pandas DataFrame.""" if not self._data: return pd.DataFrame() return self._data.table @property - def chart(self): + def chart(self) -> Any: """Returns a chart back from the aggregation.""" if not self._data: return None return self._data.chart - def from_dict(self, data_dict): - """Feed a block from a block dict.""" + def from_dict(self, data_dict: Dict[str, Any]) -> None: + """Feed a block from a block dict. + + Args: + data_dict: Dictionary with block data. + """ component = data_dict.get("componentName", "N/A") if component != "TsAggregationCompact": raise TypeError("Not an aggregation block.") @@ -299,7 +364,7 @@ def from_dict(self, data_dict): self._chart_type = agg_dict.get("chart_type", "table") self._agg_dict = agg_dict - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the block data. Raises: @@ -360,49 +425,63 @@ class AggregationGroupBlock(BaseBlock): TYPE = "aggregation_group" - def __init__(self, story, index): + def __init__(self, story: Story, index: int) -> None: + """Initialize the aggregation group block. + + Args: + story: Story object. + index: Index of the block. + """ super().__init__(story, index) self._group_id = 0 self._group_name = "" @property - def group(self): + def group(self) -> Optional[aggregation.AggregationGroup]: """Returns the aggregation group object.""" if self._data: return self._data return None @group.setter - def group(self, group_obj): - """Set the aggregation group object.""" + def group(self, group_obj: aggregation.AggregationGroup) -> None: + """Set the aggregation group object. + + Args: + group_obj: AggregationGroup object. + """ self._data = group_obj @property - def group_name(self): + def group_name(self) -> str: """Returns the aggregation group name.""" return self._group_name @property - def group_id(self): + def group_id(self) -> int: """Returns the aggregation group ID.""" return self._group_id @property - def table(self): + def table(self) -> pd.DataFrame: """Returns a table view, as a pandas DataFrame.""" if not self._data: return pd.DataFrame() return self._data.table @property - def chart(self): + def chart(self) -> Any: """Returns a chart back from the aggregation.""" if not self._data: return None return self._data.chart - def from_dict(self, data_dict): - """Feed a block from a block dict.""" + def from_dict(self, data_dict: Dict[str, Any]) -> None: + """Feed a block from a block dict. + + Args: + data_dict: Dictionary with block data. + """ component = data_dict.get("componentName", "N/A") if component != "TsAggregationGroupCompact": raise TypeError("Not an aggregation group block.") @@ -415,7 +494,7 @@ def from_dict(self, data_dict): self._group_id = group_dict.get("id", 0) self._group_name = group_dict.get("name", "") - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: """Returns a dict with the block data. Raises: @@ -452,7 +531,7 @@ class Story(resource.BaseResource): id: Primary key of the story. """ - def __init__(self, story_id, sketch, api): + def __init__(self, story_id: int, sketch: Sketch, api: TimesketchApi) -> None: """Initializes the Story object. Args: @@ -463,26 +542,26 @@ def __init__(self, story_id, sketch, api): self.id = story_id self._api = api self._title = "" - self._blocks = [] + self._blocks: List[BaseBlock] = [] self._sketch = sketch resource_uri = "sketches/{0:d}/stories/{1:d}/".format(sketch.id, self.id) super().__init__(api, resource_uri) @property - def blocks(self): + def blocks(self) -> List[BaseBlock]: """Returns all the blocks of the story.""" if not self._blocks: story_data = self.lazyload_data(refresh_cache=True) objects = story_data.get("objects") - content = "" + content = "[]" if objects: - content = objects[0].get("content", []) + content = objects[0].get("content", "[]") index = 0 for content_block in json.loads(content): name = content_block.get("componentName", "") if content_block.get("content"): - block = TextBlock(self, index) + block: BaseBlock = TextBlock(self, index) block.from_dict(content_block) elif name == "TsViewEventList": block = ViewBlock(self, index) @@ -506,12 +585,15 @@ def blocks(self): group_obj = aggregation.AggregationGroup(self._sketch) group_obj.from_saved(block.group_id) block.feed(group_obj) + else: + continue + self._blocks.append(block) index += 1 return self._blocks @property - def title(self): + def title(self) -> str: """Property that returns story title. Returns: @@ -525,7 +607,7 @@ def title(self): return self._title @property - def content(self): + def content(self) -> str: """Property that returns the content of a story. Returns: @@ -535,30 +617,41 @@ def content(self): return json.dumps(content_list) @property - def size(self): + def size(self) -> int: """Retiurns the number of blocks stored in the story.""" return len(self._blocks) - def __len__(self): + def __len__(self) -> int: """Returns the number of blocks stored in the story.""" _ = self.blocks return len(self._blocks) - def _add_block(self, block, index): - """Adds a block to the story's content.""" + def _add_block(self, block: BaseBlock, index: int) -> bool: + """Adds a block to the story's content. + + Args: + block: Story block object. + index: Index to add the block at. + """ self._blocks.insert(index, block) self.commit() self.reset() - - def add_aggregation(self, agg_obj, chart_type="table", index=-1): + return True + + def add_aggregation( + self, + agg_obj: aggregation.Aggregation, + chart_type: str = "table", + index: int = -1, + ) -> bool: """Adds an aggregation object to the story. Args: - agg_obj (aggregation.Aggregation): an aggregation object - chart_type (str): string indicating the type of aggregation, can be: + agg_obj: an aggregation object + chart_type: string indicating the type of aggregation, can be: "table" or the name of the chart to be used, eg "barcharct", "hbarchart". Defaults to "table". - index (int): an integer, if supplied determines where the new + index: an integer, if supplied determines where the new block will be added. If not supplied it will be appended at the end. @@ -583,7 +676,7 @@ def add_aggregation(self, agg_obj, chart_type="table", index=-1): return self._add_block(agg_block, index) - def add_text(self, text, index=-1): + def add_text(self, text: str, index: int = -1) -> bool: """Adds a text block to the story. Args: @@ -602,7 +695,7 @@ def add_text(self, text, index=-1): return self._add_block(text_block, index) - def add_view(self, view_obj, index=-1): + def add_view(self, view_obj: Any, index: int = -1) -> bool: """Add a view to the story. Args: @@ -617,14 +710,14 @@ def add_view(self, view_obj, index=-1): Raises: TypeError: if the view object is not of the correct type. """ - self.add_saved_search(view_obj, index) + return self.add_saved_search(view_obj, index) - def add_saved_search(self, search_obj, index=-1): + def add_saved_search(self, search_obj: search.Search, index: int = -1) -> bool: """Add a saved search to the story. Args: - search_obj (search.Search): a search object - index (int): an integer, if supplied determines where the new + search_obj: a search object + index: an integer, if supplied determines where the new block will be added. If not supplied it will be appended at the end. @@ -648,7 +741,7 @@ def add_saved_search(self, search_obj, index=-1): return self._add_block(view_block, index) - def commit(self): + def commit(self) -> bool: """Commit the story to the server.""" content_list = [x.to_dict() for x in self._blocks] content = json.dumps(content_list) @@ -663,7 +756,7 @@ def commit(self): return error.check_return_status(response, logger) - def delete(self): + def delete(self) -> bool: """Delete the story from the sketch. Returns: @@ -675,8 +768,13 @@ def delete(self): return error.check_return_status(response, logger) - def move_to(self, block, new_index): - """Moves a block from one index to another.""" + def move_to(self, block: BaseBlock, new_index: int) -> None: + """Moves a block from one index to another. + + Args: + block: Story block object. + new_index: Index to move the block to. + """ if new_index < 0: return old_index = block.index @@ -686,31 +784,39 @@ def move_to(self, block, new_index): self._blocks.insert(new_index, block) self.commit() - def remove_block(self, index): - """Removes a block from the story.""" + def remove_block(self, index: int) -> None: + """Removes a block from the story. + + Args: + index: Index of the block to remove. + """ _ = self._blocks.pop(index) self.commit() self.reset() - def reset(self): + def reset(self) -> None: """Refresh story content.""" self._title = "" self._blocks = [] _ = self.lazyload_data(refresh_cache=True) _ = self.blocks - def to_html(self): + def to_html(self) -> str: """Returns HTML formatted string with the content of the story.""" story_dict = self.to_export_format("html") - return story_dict.get("story", "") + return str(story_dict.get("story", "")) - def to_markdown(self): + def to_markdown(self) -> str: """Returns markdown formatted string with the content of the story.""" story_dict = self.to_export_format("markdown") - return story_dict.get("story", "") + return str(story_dict.get("story", "")) + + def to_export_format(self, export_format: str) -> Dict[str, Any]: + """Returns exported copy of the story as defined in export_format. - def to_export_format(self, export_format): - """Returns exported copy of the story as defined in export_format.""" + Args: + export_format: The format to export the story as. + """ resource_url = "{0:s}/sketches/{1:d}/stories/{2:d}/".format( self._api.api_root, self._sketch.id, self.id ) @@ -720,15 +826,16 @@ def to_export_format(self, export_format): return error.get_response_json(response, logger) - def to_string(self): + def to_string(self) -> str: """Returns a string with the content of all the story.""" self.reset() - string_list = [] + string_list: List[str] = [] for block in self.blocks: if block.TYPE == "text": - string_list.append(block.text) + string_list.append(str(block.data)) elif block.TYPE == "view": - search_obj = block.view + # Type casting for search_obj as it is expected to be search.Search + search_obj: Optional[search.Search] = block.view if search_obj is None: logging.warning("Block has no view. Skipping") continue diff --git a/api_client/python/timesketch_api_client/story_test.py b/api_client/python/timesketch_api_client/story_test.py index d6651080aa..2cc96d3d1e 100644 --- a/api_client/python/timesketch_api_client/story_test.py +++ b/api_client/python/timesketch_api_client/story_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import client from . import test_lib diff --git a/api_client/python/timesketch_api_client/test_lib.py b/api_client/python/timesketch_api_client/test_lib.py index eb2b51b347..d988b27e06 100644 --- a/api_client/python/timesketch_api_client/test_lib.py +++ b/api_client/python/timesketch_api_client/test_lib.py @@ -13,8 +13,6 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import json auth_text_data = '' diff --git a/api_client/python/timesketch_api_client/timeline.py b/api_client/python/timesketch_api_client/timeline.py index fb6acaee70..5191ef340f 100644 --- a/api_client/python/timesketch_api_client/timeline.py +++ b/api_client/python/timesketch_api_client/timeline.py @@ -1,4 +1,4 @@ -# Copyright 2019 Google Inc. All rights reserved. +# http://www.apache.org/licenses/LICENSE-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,16 +13,25 @@ # limitations under the License. """Timesketch API client library.""" -from __future__ import unicode_literals +from __future__ import annotations import json import logging +from typing import Any +from typing import Dict +from typing import List +from typing import Optional +from typing import TYPE_CHECKING + from . import analyzer from . import error from . import index from . import resource +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.timeline") @@ -33,7 +42,14 @@ class Timeline(resource.BaseResource): id: Primary key of the timeline. """ - def __init__(self, timeline_id, sketch_id, api, name=None, searchindex=None): + def __init__( + self, + timeline_id: int, + sketch_id: int, + api: TimesketchApi, + name: Optional[str] = None, + searchindex: Optional[str] = None, + ) -> None: """Initializes the Timeline object. Args: @@ -53,7 +69,7 @@ def __init__(self, timeline_id, sketch_id, api, name=None, searchindex=None): super().__init__(api, resource_uri) @property - def labels(self): + def labels(self) -> List[str]: """Property that returns the timeline labels.""" data = self.lazyload_data(refresh_cache=True) objects = data.get("objects", []) @@ -68,7 +84,7 @@ def labels(self): return [] @property - def color(self): + def color(self) -> str: """Property that returns timeline color. Returns: @@ -80,13 +96,17 @@ def color(self): return self._color @color.setter - def color(self, color): - """Change the color of the timeline.""" + def color(self, color: str) -> None: + """Change the color of the timeline. + + Args: + color: The new color of the timeline. + """ self._color = color self._commit() @property - def data_sources(self): + def data_sources(self) -> List[Dict[str, Any]]: """Property that returns the timeline data sources.""" data = self.lazyload_data(refresh_cache=True) objects = data.get("objects", []) @@ -97,7 +117,7 @@ def data_sources(self): return timeline_data.get("datasources", []) @property - def description(self): + def description(self) -> str: """Property that returns timeline description. Returns: @@ -109,13 +129,17 @@ def description(self): return self._description @description.setter - def description(self, description): - """Change the timeline description.""" + def description(self, description: str) -> None: + """Change the timeline description. + + Args: + description: The new description of the timeline. + """ self._description = description self._commit() @property - def name(self): + def name(self) -> str: """Property that returns timeline name. Returns: @@ -127,13 +151,17 @@ def name(self): return self._name @name.setter - def name(self, name): - """Change the name of the timeline.""" + def name(self, name: str) -> None: + """Change the name of the timeline. + + Args: + name: The new name of the timeline. + """ self._name = name self._commit() @property - def index(self): + def index(self) -> Optional[index.SearchIndex]: """Property that returns index object. Returns: @@ -153,7 +181,7 @@ def index(self): ) @property - def index_name(self): + def index_name(self) -> str: """Property that returns index name. Returns: @@ -165,7 +193,7 @@ def index_name(self): self._searchindex = index_name return self._searchindex - def is_archived(self): + def is_archived(self) -> bool: """Return a boolean indicating whether the timeline is archived.""" resource_url = f"{self.api.api_root}/sketches/{self._sketch_id}/archive/" response = self.api.session.get(resource_url) @@ -178,16 +206,21 @@ def is_archived(self): return sketch_is_archived return timeline_dict.get(self.index_name) - def run_analyzer(self, analyzer_name, analyzer_kwargs=None, ignore_previous=False): + def run_analyzer( + self, + analyzer_name: str, + analyzer_kwargs: Optional[Dict[str, Any]] = None, + ignore_previous: bool = False, + ) -> List[analyzer.AnalyzerResult]: """Run an analyzer on a timeline. Args: - analyzer_name (str): a name of an analyzer class to run against the + analyzer_name: a name of an analyzer class to run against the timeline. - analyzer_kwargs (dict): optional dict with parameters for the analyzer. + analyzer_kwargs: optional dict with parameters for the analyzer. This is optional and just for those analyzers that can accept further parameters. - ignore_previous (bool): an optional bool, if set to True then + ignore_previous: an optional bool, if set to True then analyzer is run irrelevant on whether it has been previously been run. @@ -223,19 +256,22 @@ def run_analyzer(self, analyzer_name, analyzer_kwargs=None, ignore_previous=Fals ) def run_analyzers( - self, analyzer_names, analyzer_kwargs=None, ignore_previous=False - ): + self, + analyzer_names: List[str], + analyzer_kwargs: Optional[Dict[str, Any]] = None, + ignore_previous: bool = False, + ) -> List[analyzer.AnalyzerResult]: """Run an analyzer on a timeline. Args: - analyzer_names (list): a list of analyzer class names to run against the + analyzer_names: a list of analyzer class names to run against the timeline. - analyzer_kwargs (dict): optional dict with parameters for the analyzer. + analyzer_kwargs: optional dict with parameters for the analyzer. This is optional and just for those analyzers that can accept further parameters. It is expected that this is a dict with the key value being the analyzer name, and the value being another key/value dict with the parameters for that analyzer. - ignore_previous (bool): an optional bool, if set to True then + ignore_previous: an optional bool, if set to True then analyzer is run irrelevant on whether it has been previously been run. @@ -278,7 +314,7 @@ def run_analyzers( "'ignore_previous=True' to overwrite.", analyzer_names, ) - return None + return [] analyzer_results = [] for session_dict in objects[0]: @@ -303,7 +339,7 @@ def run_analyzers( return analyzer_results @property - def status(self): + def status(self) -> str: """Property that returns the timeline status. Returns: @@ -319,7 +355,7 @@ def status(self): status = status_list[0] return status.get("status") - def _commit(self): + def _commit(self) -> bool: """Commit changes to the timeline.""" resource_url = "{0:s}/{1:s}".format(self.api.api_root, self.resource_uri) @@ -336,14 +372,14 @@ def _commit(self): return status - def add_timeline_label(self, label): + def add_timeline_label(self, label: str) -> bool: """Add a label to the timeline. Args: - label (str): A string with the label to add to the timeline. + label: A string with the label to add to the timeline. Returns: - bool: A boolean to indicate whether the label was successfully + A boolean to indicate whether the label was successfully added to the timeline. """ if label in self.labels: @@ -367,14 +403,14 @@ def add_timeline_label(self, label): return status - def remove_timeline_label(self, label): + def remove_timeline_label(self, label: str) -> bool: """Remove a label from the timeline. Args: - label (str): A string with the label to remove from the timeline. + label: A string with the label to remove from the timeline. Returns: - bool: A boolean to indicate whether the label was successfully + A boolean to indicate whether the label was successfully removed from the timeline. """ if label not in self.labels: @@ -401,7 +437,7 @@ def remove_timeline_label(self, label): return status - def delete(self): + def delete(self) -> bool: """Deletes the timeline.""" resource_url = "{0:s}/{1:s}".format(self.api.api_root, self.resource_uri) response = self.api.session.delete(resource_url) diff --git a/api_client/python/timesketch_api_client/timeline_test.py b/api_client/python/timesketch_api_client/timeline_test.py index cc223e78cc..af82605943 100644 --- a/api_client/python/timesketch_api_client/timeline_test.py +++ b/api_client/python/timesketch_api_client/timeline_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import client from . import test_lib diff --git a/api_client/python/timesketch_api_client/user.py b/api_client/python/timesketch_api_client/user.py index d7833675dd..aa6e4143fa 100644 --- a/api_client/python/timesketch_api_client/user.py +++ b/api_client/python/timesketch_api_client/user.py @@ -13,20 +13,31 @@ # limitations under the License. """Timesketch API client library.""" +from __future__ import annotations + import logging +from typing import Any, Dict, List, Optional, TYPE_CHECKING from . import error from . import resource +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.user") class User(resource.BaseResource): """User object.""" - def __init__(self, api, user_id=None): - """Initializes the user object.""" - self._object_data = None + def __init__(self, api: TimesketchApi, user_id: Optional[int] = None) -> None: + """Initializes the user object. + + Args: + api: An instance of TimesketchApi object. + user_id: Primary key ID of the user (optional). + """ + self._object_data: Optional[Dict[str, Any]] = None if not user_id: resource_uri = "users/me/" super().__init__(api, resource_uri) @@ -35,7 +46,7 @@ def __init__(self, api, user_id=None): self.api = api super().__init__(api=api, resource_uri=f"users/{self.id}") - def _get_data(self): + def _get_data(self) -> Dict[str, Any]: """Returns dict from the first object of the resource data.""" if self._object_data: return self._object_data @@ -49,17 +60,17 @@ def _get_data(self): return self._object_data - def change_password(self, new_password): + def change_password(self, new_password: str) -> bool: """Change the password for the user. Args: - new_password (str): String with the password. + new_password: String with the password. Raises: ValueError: If there was an error. Returns: - Boolean: Whether the password was successfully modified. + Whether the password was successfully modified. """ if not new_password: raise ValueError("No new password supplied.") @@ -73,31 +84,31 @@ def change_password(self, new_password): return error.check_return_status(response, logger) @property - def groups(self): + def groups(self) -> List[str]: """Property that returns the groups the user belongs to.""" data = self._get_data() groups = data.get("groups", []) - return [x.get("name", "") for x in groups] + return [str(x.get("name", "")) for x in groups] @property - def is_active(self): + def is_active(self) -> bool: """Property that returns bool indicating whether the user is active.""" data = self._get_data() - return data.get("active", True) + return bool(data.get("active", True)) @property - def is_admin(self): + def is_admin(self) -> bool: """Property that returns bool indicating whether the user is admin.""" data = self._get_data() - return data.get("admin", False) + return bool(data.get("admin", False)) @property - def username(self): + def username(self) -> str: """Property that returns back the username of the current user.""" data = self._get_data() - return data.get("username", "Unknown") + return str(data.get("username", "Unknown")) - def __str__(self): + def __str__(self) -> str: """Returns a string representation of the username.""" user_strings = [self.username] diff --git a/api_client/python/timesketch_api_client/view.py b/api_client/python/timesketch_api_client/view.py index 5c4574c912..7d019d6c20 100644 --- a/api_client/python/timesketch_api_client/view.py +++ b/api_client/python/timesketch_api_client/view.py @@ -13,13 +13,17 @@ # limitations under the License. """Timesketch API client library.""" -from __future__ import unicode_literals +from __future__ import annotations import json import logging +from typing import Any, Dict, Union, TYPE_CHECKING from . import resource +if TYPE_CHECKING: + from .client import TimesketchApi + logger = logging.getLogger("timesketch_api.view") @@ -31,7 +35,9 @@ class View(resource.BaseResource): name: Name of the view. """ - def __init__(self, view_id, view_name, sketch_id, api): + def __init__( + self, view_id: int, view_name: str, sketch_id: int, api: TimesketchApi + ) -> None: """Initializes the View object. Args: @@ -49,7 +55,9 @@ def __init__(self, view_id, view_name, sketch_id, api): resource_uri = "sketches/{0:d}/views/{1:d}/".format(sketch_id, self.id) super().__init__(api, resource_uri) - def _get_top_level_attribute(self, name, default_value=None, refresh=False): + def _get_top_level_attribute( + self, name: str, default_value: Any = None, refresh: bool = False + ) -> Any: """Returns a top level attribute from a view object. Args: @@ -59,7 +67,7 @@ def _get_top_level_attribute(self, name, default_value=None, refresh=False): refresh: If set to True then the data will be refreshed. Returns: - The dict value of the key "name". + The value of the key "name". """ view = self.lazyload_data(refresh_cache=refresh) view_objects = view.get("objects") @@ -72,29 +80,29 @@ def _get_top_level_attribute(self, name, default_value=None, refresh=False): return first_object.get(name, default_value) @property - def description(self): + def description(self) -> str: """Property that returns the description value of a view. Returns: - Description of the view as a string. + Description of the view. """ return self._get_top_level_attribute("description", default_value="") @property - def user(self): + def user(self) -> str: """Property that returns the username of the view creator. Returns: - A string with the username of the user generating the view. + The username of the user generating the view. """ user_dict = self._get_top_level_attribute("user", default_value={}) username = user_dict.get("username") if not username: return "System" - return username + return str(username) @property - def query_string(self): + def query_string(self) -> str: """Property that returns the views query string. Returns: @@ -103,7 +111,7 @@ def query_string(self): return self._get_top_level_attribute("query_string", default_value="") @property - def query_filter(self): + def query_filter(self) -> Union[Dict[str, Any], str]: """Property that returns the views filter. Returns: @@ -117,7 +125,7 @@ def query_filter(self): return json.loads(query_filter_string) @property - def query_dsl(self): + def query_dsl(self) -> Union[Dict[str, Any], str]: """Property that returns the views query DSL. Returns: diff --git a/api_client/python/timesketch_api_client/view_test.py b/api_client/python/timesketch_api_client/view_test.py index df20f6e6d5..69461672ce 100644 --- a/api_client/python/timesketch_api_client/view_test.py +++ b/api_client/python/timesketch_api_client/view_test.py @@ -13,10 +13,8 @@ # limitations under the License. """Tests for the Timesketch API client""" -from __future__ import unicode_literals - import unittest -import mock +from unittest import mock from . import client from . import search diff --git a/importer_client/python/timesketch_import_client/importer.py b/importer_client/python/timesketch_import_client/importer.py index 552fd81659..1919289687 100644 --- a/importer_client/python/timesketch_import_client/importer.py +++ b/importer_client/python/timesketch_import_client/importer.py @@ -1036,7 +1036,7 @@ def set_max_payload_size(self, size_in_bytes: int) -> None: """Set the maximum payload size allowed by the server. Args: - size_in_bytes (int): The server limit (e.g. MAX_FORM_MEMORY_SIZE). + size_in_bytes: The server limit (e.g. MAX_FORM_MEMORY_SIZE). """ if size_in_bytes <= 0: raise ValueError(f"Payload size must be positive, got {size_in_bytes}") diff --git a/importer_client/python/tools/timesketch_importer.py b/importer_client/python/tools/timesketch_importer.py index 82981b51dd..60aa6198d6 100644 --- a/importer_client/python/tools/timesketch_importer.py +++ b/importer_client/python/tools/timesketch_importer.py @@ -65,10 +65,10 @@ def upload_file( """Uploads a file to Timesketch. Args: - my_sketch (sketch.Sketch): a sketch object to point to the sketch the + my_sketch: a sketch object to point to the sketch the data will be imported to. - config_dict (dict): dict with settings for the importer. - file_path (str): the path to the file to upload. + config_dict: dict with settings for the importer. + file_path: the path to the file to upload. Returns: A tuple with the timeline object (timeline.Timeline) or None if not