Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update Turbinia dependencies #793

Merged
merged 6 commits into from
Oct 23, 2023
Merged

Update Turbinia dependencies #793

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
b239f63
Update dependencies
jleaniz Oct 20, 2023
a31b6c7
Update type annotations
jleaniz Oct 20, 2023
805b430
Update types and annotations
jleaniz Oct 20, 2023
abd5ce4
Updates
jleaniz Oct 20, 2023
5866495
Appease linter
jleaniz Oct 20, 2023
333880e
fix file formatting issue
jleaniz Oct 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 22 additions & 17 deletions dftimewolf/lib/processors/turbinia_base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,13 +25,13 @@
# pylint: disable=unused-import
from dftimewolf.lib import state as state_lib


# mypy: disable-error-code="attr-defined"
# pylint: disable=abstract-method,no-member
class TurbiniaProcessorBase(module.BaseModule):
"""Base class for processing with Turbinia.

Attributes:
client (turbinia_api_lib.ApiClient): Turbinia client.
client (turbinia_api_lib.api_client.ApiClient): Turbinia client.
client_config (dict): Turbinia client config.
credentials google.oauth2.credentials.Credentials: User Oauth2 credentials.
extensions (List[str]): List of file extensions to look for.
Expand Down Expand Up @@ -69,8 +69,9 @@ def __init__(
the entire recipe to fail if the module encounters an error.
"""
super().__init__(state=state, name=name, critical=critical)
self.client: Optional[turbinia_api_lib.ApiClient] = None
self.client_config: Optional[turbinia_api_lib.Configuration] = None
self.client: Optional[turbinia_api_lib.api_client.ApiClient] = None
# pylint: disable=line-too-long
self.client_config: Optional[turbinia_api_lib.configuration.Configuration] = None
self.client_secrets_path = os.path.join(
os.path.expanduser('~'), ".dftimewolf_turbinia_secrets.json")
self.credentials_path = os.path.join(
Expand All @@ -88,7 +89,7 @@ def __init__(
self.output_path = str()
self.parallel_count = 5 # Arbitrary, used by ThreadAwareModule
self.project = str()
self.requests_api_instance: turbinia_requests_api.TurbiniaRequestsApi = None
self.requests_api_instance: turbinia_requests_api.TurbiniaRequestsApi = None # type: ignore
self.sketch_id = int()
self.state = state
self.turbinia_auth = bool()
Expand Down Expand Up @@ -171,7 +172,7 @@ def DownloadFilesFromAPI(self, task_data: Dict[str, List[str]],
try:
task_id = task_data.get('id')
api_response = api_instance.get_task_output_with_http_info(
task_id, _preload_content=False)
task_id, _preload_content=False) # type: ignore
filename = f'{task_id}-'

# Create a temporary file to write the response to.
Expand All @@ -180,7 +181,8 @@ def DownloadFilesFromAPI(self, task_data: Dict[str, List[str]],
local_path = file.name
self.logger.info(f'Downloading output for task {task_id} to {local_path}')
# Read the response and write to the file.
file.write(api_response.raw_data)
if api_response.raw_data:
file.write(api_response.raw_data)
file.close()

# Extract the files from the tgz file.
Expand All @@ -189,7 +191,7 @@ def DownloadFilesFromAPI(self, task_data: Dict[str, List[str]],
result = extracted_path
self.PublishMessage(
f'Extracted output file to {result} for task {task_id}')
except turbinia_api_lib.ApiException as exception:
except turbinia_api_lib.exceptions.ApiException as exception:
self.ModuleError(
f'Unable to download task data: {exception}', critical=False)
except OSError as exception:
Expand Down Expand Up @@ -254,15 +256,16 @@ def RefreshClientCredentials(self) -> bool:
return bool(refresh)

def InitializeTurbiniaApiClient(
self, credentials: Credentials) -> turbinia_api_lib.ApiClient:
self, credentials: Credentials) -> turbinia_api_lib.api_client.ApiClient:
"""Creates a Turbinia API client object.

This method also runs the authentication flow if needed.

Returns:
turbinia_api_lib.ApiClient: A Turbinia API client object.
turbinia_api_lib.api_client.ApiClient: A Turbinia API client object.
"""
self.client_config = turbinia_api_lib.Configuration(host=self.turbinia_api)
self.client_config = turbinia_api_lib.configuration.Configuration(
host=self.turbinia_api)
if not self.client_config:
self.ModuleError('Unable to configure Turbinia API server', critical=True)
# Check if Turbinia requires authentication.
Expand All @@ -275,7 +278,7 @@ def InitializeTurbiniaApiClient(
else:
self.ModuleError(
'Unable to obtain id_token from identity provider', critical=True)
return turbinia_api_lib.ApiClient(self.client_config)
return turbinia_api_lib.api_client.ApiClient(self.client_config)

def TurbiniaSetUp(
self, project: str, turbinia_auth: bool,
Expand All @@ -299,7 +302,8 @@ def TurbiniaSetUp(
self.turbinia_zone = turbinia_zone
self.incident_id = incident_id
self.sketch_id = sketch_id
self.client_config = turbinia_api_lib.Configuration(host=self.turbinia_api)
self.client_config = turbinia_api_lib.configuration.Configuration(
host=self.turbinia_api)
self.client = self.InitializeTurbiniaApiClient(self.credentials)
self.requests_api_instance = turbinia_requests_api.TurbiniaRequestsApi(
self.client)
Expand All @@ -309,7 +313,7 @@ def TurbiniaSetUp(
try:
api_response = api_instance.read_config()
self.output_path = api_response.get('OUTPUT_DIR')
except turbinia_api_lib.ApiException as exception:
except turbinia_api_lib.exceptions.ApiException as exception:
self.ModuleError(exception.body, critical=True)

def TurbiniaStart(
Expand Down Expand Up @@ -367,14 +371,15 @@ def TurbiniaStart(

# Send the request to the API server.
try:
response = self.requests_api_instance.create_request(request)
response = self.requests_api_instance.create_request(
request) # type: ignore
request_id = response.get('request_id')
self.PublishMessage(
'Creating Turbinia request {0!s} with evidence {1!s}'.format(
request_id, evidence_name))
self.PublishMessage(
'Turbinia request status at {0!s}'.format(self.turbinia_api))
except turbinia_api_lib.ApiException as exception:
except turbinia_api_lib.exceptions.ApiException as exception:
self.ModuleError(str(exception), critical=True)

if not request_id:
Expand Down Expand Up @@ -441,7 +446,7 @@ def TurbiniaWait(self,
processed_paths.add(path)
yield task, path

except turbinia_api_lib.ApiException as exception:
except turbinia_api_lib.exceptions.ApiException as exception:
retries += 1
self.logger.warning(f'Retrying after exception: {exception.body}')

Expand Down
Loading
Loading