Skip to content

Commit

Permalink
black and isort linter fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
mandar242 committed Feb 27, 2025
1 parent c9c3cc3 commit 657ca6a
Showing 1 changed file with 44 additions and 14 deletions.
58 changes: 44 additions & 14 deletions plugins/connection/aws_ssm.py
Original file line number Diff line number Diff line change
Expand Up @@ -327,13 +327,13 @@
import string
import subprocess
import time
from typing import Optional
from typing import NoReturn
from typing import Tuple
from typing import Any
from typing import List
from typing import Dict
from typing import Iterator
from typing import List
from typing import NoReturn
from typing import Optional
from typing import Tuple

try:
import boto3
Expand Down Expand Up @@ -491,7 +491,13 @@ def get_bucket_endpoint(self) -> Tuple[str, str]:

return s3_bucket_client.meta.endpoint_url, s3_bucket_client.meta.region_name

def get_boto_client(self, service: str, region_name: Optional[str] = None, profile_name: Optional[str] = None, endpoint_url: Optional[str] = None) -> Any:
def get_boto_client(
self,
service: str,
region_name: Optional[str] = None,
profile_name: Optional[str] = None,
endpoint_url: Optional[str] = None,
) -> Any:
"""Gets a boto3 client based on the STS token"""

aws_access_key_id = self.connection.get_option("access_key_id")
Expand All @@ -518,7 +524,14 @@ def get_boto_client(self, service: str, region_name: Optional[str] = None, profi
)
return client

def get_url(self, client_method: str, bucket_name: str, out_path: str, http_method: str, extra_args: Optional[Dict[str, Any]] = None) -> str:
def get_url(
self,
client_method: str,
bucket_name: str,
out_path: str,
http_method: str,
extra_args: Optional[Dict[str, Any]] = None,
) -> str:
"""Generate URL for get_object / put_object"""

client = self._s3_client
Expand All @@ -536,9 +549,13 @@ def generate_encryption_settings(self) -> Tuple[Dict, Dict]:

put_args["ServerSideEncryption"] = self.connection.get_option("bucket_sse_mode")
put_headers["x-amz-server-side-encryption"] = self.connection.get_option("bucket_sse_mode")
if self.connection.get_option("bucket_sse_mode") == "aws:kms" and self.connection.get_option("bucket_sse_kms_key_id"):
if self.connection.get_option("bucket_sse_mode") == "aws:kms" and self.connection.get_option(
"bucket_sse_kms_key_id"
):
put_args["SSEKMSKeyId"] = self.connection.get_option("bucket_sse_kms_key_id")
put_headers["x-amz-server-side-encryption-aws-kms-key-id"] = self.connection.get_option("bucket_sse_kms_key_id")
put_headers["x-amz-server-side-encryption-aws-kms-key-id"] = self.connection.get_option(
"bucket_sse_kms_key_id"
)
return put_args, put_headers


Expand Down Expand Up @@ -644,9 +661,7 @@ def _initialize_s3_client(self, profile_name: str) -> None:
self._vvvv(f"SETUP BOTO3 CLIENTS: S3 {s3_endpoint_url}")

self.s3_manager.initialize_s3_client(
region_name=s3_region_name,
endpoint_url=s3_endpoint_url,
profile_name=profile_name
region_name=s3_region_name, endpoint_url=s3_endpoint_url, profile_name=profile_name
)

def _display(self, f: Any, message: str) -> None:
Expand Down Expand Up @@ -976,11 +991,24 @@ def _flush_stderr(self, session_process) -> str:

return stderr

def _get_url(self, client_method: str, bucket_name: str, out_path: str, http_method: str, extra_args: Optional[Dict[str, Any]] = None) -> str:
def _get_url(
self,
client_method: str,
bucket_name: str,
out_path: str,
http_method: str,
extra_args: Optional[Dict[str, Any]] = None,
) -> str:
"""Generate URL for get_object / put_object"""
return self.s3_manager.get_url(client_method, bucket_name, out_path, http_method, extra_args)

def _get_boto_client(self, service: str, region_name: Optional[str] = None, profile_name: Optional[str] = None, endpoint_url: Optional[str] = None):
def _get_boto_client(
self,
service: str,
region_name: Optional[str] = None,
profile_name: Optional[str] = None,
endpoint_url: Optional[str] = None,
):
"""Gets a boto3 client based on the STS token"""
return self.s3_manager.get_boto_client(service, region_name, profile_name, endpoint_url)

Expand All @@ -990,7 +1018,9 @@ def _escape_path(self, path: str) -> str:
def _generate_encryption_settings(self) -> Tuple[Dict, Dict]:
return self.s3_manager.generate_encryption_settings()

def _generate_commands(self, bucket_name: str, s3_path: str, in_path: str, out_path: str) -> Tuple[List, List, Dict]:
def _generate_commands(
self, bucket_name: str, s3_path: str, in_path: str, out_path: str
) -> Tuple[List, List, Dict]:
put_args, put_headers = self._generate_encryption_settings()

put_url = self._get_url("put_object", bucket_name, s3_path, "PUT", extra_args=put_args)
Expand Down

0 comments on commit 657ca6a

Please sign in to comment.