adding delete and put interfaces for the S3 plugin

This commit is contained in:
Hossein Shafagh 2020-10-15 18:13:50 -07:00
parent cd29b2b870
commit bfe89e131e
2 changed files with 76 additions and 2 deletions

View File

@ -9,9 +9,12 @@
from flask import current_app
from .sts import sts_client
from botocore.exceptions import ClientError
from lemur.extensions import sentry
@sts_client("s3", service_type="resource")
def put(bucket_name, region, prefix, data, encrypt, **kwargs):
def put(bucket_name, prefix, data, encrypt, **kwargs):
"""
Use STS to write to an S3 bucket
"""
@ -32,4 +35,37 @@ def put(bucket_name, region, prefix, data, encrypt, **kwargs):
ServerSideEncryption="AES256",
)
else:
bucket.put_object(Key=prefix, Body=data, ACL="bucket-owner-full-control")
try:
bucket.put_object(Key=prefix, Body=data, ACL="bucket-owner-full-control")
except ClientError:
sentry.captureException()
@sts_client("s3", service_type="client")
def delete(bucket_name, prefix, **kwargs):
"""
Use STS to delete an object
"""
try:
response = kwargs["client"].delete_object(Bucket=bucket_name, Key=prefix)
current_app.logger.debug(f"Delete data from S3."
f"Bucket: {bucket_name},"
f"Prefix: {prefix},"
f"Status_code: {response}")
return response['ResponseMetadata']['HTTPStatusCode'] < 300
except ClientError:
sentry.captureException()
@sts_client("s3", service_type="client")
def get(bucket_name, prefix, **kwargs):
"""
Use STS to get an object
"""
try:
response = kwargs["client"].get_object(Bucket=bucket_name, Key=prefix)
current_app.logger.debug(f"Get data from S3. Bucket: {bucket_name},"
f"Prefix: {prefix}")
return response['Body'].read().decode("utf-8")
except ClientError:
sentry.captureException()

View File

@ -0,0 +1,38 @@
import boto3
from moto import mock_sts, mock_s3
@mock_sts()
@mock_s3()
def test_put_delete_s3_object(app):
from lemur.plugins.lemur_aws.s3 import put, delete, get
bucket = "public-bucket"
account = "123456789012"
path = "some_path/foo"
s3_client = boto3.client('s3')
s3_client.create_bucket(Bucket=bucket)
data = "dummy data"
put(bucket_name=bucket,
prefix=path,
data=data,
encrypt=None,
account_number=account)
response = get(bucket_name=bucket, prefix=path, account_number=account)
# put data, and getting the same data
assert (response == data)
response = get(bucket_name="wrong-bucket", prefix=path, account_number=account)
# attempting to get thccle wrong data
assert (response is None)
delete(bucket_name=bucket, prefix=path, account_number=account)
response = get(bucket_name=bucket, prefix=path, account_number=account)
# delete data, and getting the same data
assert (response is None)