importosimportrefromflaskimportFlask,abort,requestimportstoreGUID_RE=re.compile(r"\A[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}\Z")app=Flask(__name__)app.config["MAX_CONTENT_LENGTH"]=512filestore=store.S3Store()# Uncomment the following line for simpler local testing of this service# filestore = store.LocalStore()@app.route("/files/",methods=["POST"])defadd_file():ifrequest.headers.get("Content-Type")!="text/plain":abort(422)guid=request.headers.get("X-guid","")ifnotGUID_RE.match(guid):abort(422)filestore.save(guid,request.data)return"",201@app.route("/files/<guid>",methods=["GET"])defget_file(guid):ifnotGUID_RE.match(guid):abort(422)try:returnfilestore.read(guid),{"Content-Type":"text/plain"}exceptstore.NotFound:abort(404)@app.route("/",methods=["GET"])defroot():return"",204
stroe.py
"""Provides two instances of a filestore.
There is not intended to be any vulnerability contained within this code. This file is provided to
make it easier to test locally without needing access to an S3 bucket.
-OOO
"""
import os
import boto3
import botocore
class NotFound(Exception):
pass
class LocalStore:
def __init__(self):
import tempfile
self.upload_directory = tempfile.mkdtemp()
def read(self, key):
filepath = os.path.join(self.upload_directory, key)
try:
with open(filepath, "rb") as fp:
return fp.read()
except FileNotFoundError:
raise NotFound
def save(self, key, data):
with open(os.path.join(self.upload_directory, key), "wb") as fp:
fp.write(data)
class S3Store:
"""Credentials grant access only to resource s3://BUCKET/* and only for:
* GetObject
* PutObject
"""
def __init__(self):
self.bucket = os.environ["BUCKET"]
self.s3 = boto3.client("s3")
def read(self, key):
try:
response = self.s3.get_object(Bucket=self.bucket, Key=key)
except botocore.exceptions.ClientError as exception:
if exception.response["ResponseMetadata"]["HTTPStatusCode"] == 403:
raise NotFound
# No other exceptions encountered during testing
return response["Body"].read()
def save(self, key, data):
self.s3.put_object(
Body=data, Bucket=self.bucket, ContentType="text/plain", Key=key
)