diff --git a/app/endpoints/issues.py b/app/endpoints/issues.py index bc204f1..841d1d0 100644 --- a/app/endpoints/issues.py +++ b/app/endpoints/issues.py @@ -50,8 +50,9 @@ def new_issue(): issue_number = rv.json().get('number') # We use the newly created issue number to help name the file # upload in the s3 bucket. - image_data = get_screenshot(screenshot) - ss_uri = upload_filedata(image_data, issue_number) + image_data, content_type = get_screenshot(screenshot) + ss_uri = upload_filedata( + image_data, issue_number, content_type) rv = add_comment(ss_uri, issue_number) return ('Issue created, screenshot uploaded.', rv.status_code) return ('Issue created (without screenshot).', 201) diff --git a/app/helpers.py b/app/helpers.py index 5baf8f7..f8821ff 100644 --- a/app/helpers.py +++ b/app/helpers.py @@ -19,7 +19,9 @@ from config import S3_LOCATION from config import S3_SECRET from flaskapp import app -DATA_URI_PREFIX = 'data:image/jpeg;base64,' +DATA_URI_PREFIX = 'data:image/' +JPEG_PREFIX = 'data:image/jpeg;base64,' +PNG_PREFIX = 'data:image/png;base64,' HEADERS = { 'Authorization': 'token {0}'.format(app.config['OAUTH_TOKEN']), @@ -53,12 +55,15 @@ def add_comment(screenshot_uri, issue_number): return requests.post(uri, data=json.dumps(body), headers=HEADERS) -def upload_filedata(imagedata, issue_number): - filename = "issue-{}-screenshot.jpg".format(issue_number) +def upload_filedata(imagedata, issue_number, content_type): + file_ext = 'png' + if 'jpeg' in content_type: + file_ext = 'jpeg' + filename = "issue-{}-screenshot.{}".format(issue_number, file_ext) try: s3.upload_fileobj( imagedata, S3_BUCKET, filename, - ExtraArgs={'ACL': 'private', 'ContentType': 'image/jpeg'} + ExtraArgs={'ACL': 'private', 'ContentType': content_type} ) except Exception as e: print('Error uploading image to s3: ', e) @@ -80,7 +85,7 @@ def valid_issue_request(body, title): def has_valid_screenshot(imagedata): """Determine if the screenshot is a base64 JPEG.""" - if imagedata and DATA_URI_PREFIX in imagedata: + if imagedata and any([x in imagedata for x in (PNG_PREFIX, JPEG_PREFIX)]): return True return False @@ -91,9 +96,15 @@ def get_screenshot(imagedata): Return the base64 string without the prefix. Otherwise, return None. """ try: - imagedata = re.sub(DATA_URI_PREFIX, '', imagedata) + content_type = '' + if JPEG_PREFIX in imagedata: + imagedata = re.sub(JPEG_PREFIX, '', imagedata) + content_type = 'image/jpeg' + elif PNG_PREFIX in imagedata: + imagedata = re.sub(PNG_PREFIX, '', imagedata) + content_type = 'image/png' bindata = BytesIO(b64decode(imagedata)) - return bindata + return (bindata, content_type) except Exception as e: print('Error decoding screenshot data', e) - return None + return (None, None)