Skip to content

Commit cb6e6f1

Browse files
authored
Feature/image upload (#5)
* Started on image upload support * Switched to put * Fixed issue with request not sending file name and updated tests * Switched to auto parsing of mime_type * Implemented auth for uploading to S3 * Added session token * Added passing of folder name * Fixed inconsistent casing for BMP mime type in image upload
1 parent 3064612 commit cb6e6f1

File tree

4 files changed

+753
-1
lines changed

4 files changed

+753
-1
lines changed
Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
import hashlib
2+
3+
import requests
4+
import uuid
5+
import os
6+
from io import BytesIO
7+
from typing import Dict, List, Optional
8+
from PIL import Image
9+
from dotenv import load_dotenv
10+
11+
from botocore.auth import SigV4Auth
12+
from botocore.awsrequest import AWSRequest
13+
from botocore.credentials import Credentials
14+
15+
load_dotenv()
16+
17+
MIME_TO_FORMAT: Dict[str, List[str]] = {
18+
'image/jpeg': ['JPEG', 'JPG'],
19+
'image/png': ['PNG'],
20+
'image/gif': ['GIF'],
21+
'image/bmp': ['BMP'],
22+
}
23+
24+
FORMAT_TO_MIME: Dict[str, str] = {
25+
'JPEG': 'image/jpeg',
26+
'JPG': 'image/jpeg',
27+
'PNG': 'image/png',
28+
'GIF': 'image/gif',
29+
"BMP": 'image/bmp'
30+
}
31+
32+
class ImageUploadError(Exception):
33+
"""Custom exception for image upload failures"""
34+
pass
35+
36+
37+
class InvalidMimeTypeError(ImageUploadError):
38+
"""Exception for invalid MIME type"""
39+
pass
40+
41+
42+
class MissingEnvironmentVariableError(ImageUploadError):
43+
"""Exception for missing environment variables"""
44+
pass
45+
46+
47+
def generate_file_name(img: Image.Image) -> str:
48+
"""Generate filename for the image
49+
50+
Args:
51+
img: PIL Image object
52+
53+
Returns:
54+
Generated filename string
55+
"""
56+
unique_id: str = str(uuid.uuid4())
57+
format_ext: str = img.format.lower() if img.format else 'png'
58+
return f"{unique_id}.{format_ext}"
59+
60+
def get_s3_bucket_uri() -> str:
61+
"""Get S3 bucket URI from environment variable"""
62+
s3_uri: Optional[str] = os.getenv('S3_BUCKET_URI')
63+
64+
if not s3_uri:
65+
raise MissingEnvironmentVariableError(
66+
"S3_BUCKET_URI environment variable is not set"
67+
)
68+
69+
return s3_uri
70+
71+
72+
def get_aws_signed_request(full_url, buffer, mime_type):
73+
credentials = Credentials(
74+
access_key=os.environ['AWS_ACCESS_KEY_ID'],
75+
secret_key=os.environ['AWS_SECRET_ACCESS_KEY'],
76+
token=os.environ.get('AWS_SESSION_TOKEN', None)
77+
)
78+
79+
if hasattr(buffer, 'read'):
80+
# It's a file-like object (BytesIO, etc.)
81+
current_pos = buffer.tell() # Save current position
82+
buffer.seek(0) # Go to start
83+
data = buffer.read() # Read all data
84+
buffer.seek(current_pos) # Restore position
85+
else:
86+
# It's already bytes
87+
data = buffer
88+
89+
# Calculate content hash and length
90+
content_hash = hashlib.sha256(data).hexdigest()
91+
content_length = len(data)
92+
93+
# Create the request for signing with required headers
94+
headers = {
95+
'Content-Type': mime_type,
96+
'Content-Length': str(content_length),
97+
'x-amz-content-sha256': content_hash
98+
}
99+
100+
# Create the request for signing
101+
aws_request = AWSRequest(
102+
method='PUT',
103+
url=full_url,
104+
data=buffer,
105+
headers=headers
106+
)
107+
108+
region = os.environ.get('AWS_REGION', 'eu-west-2')
109+
110+
# Sign the request
111+
SigV4Auth(credentials, 's3', region).add_auth(aws_request)
112+
113+
return aws_request
114+
115+
116+
def upload_image(img: Image.Image, folder_name: str) -> str:
117+
"""Upload PIL image with comprehensive MIME type validation
118+
119+
Args:
120+
folder_name: name of folder to save image
121+
img: PIL Image object to upload
122+
123+
Returns:
124+
JSON response from the server as a dictionary
125+
126+
Raises:
127+
InvalidMimeTypeError: If MIME type validation fails
128+
MissingEnvironmentVariableError: If S3_BUCKET_URI is not set
129+
ImageUploadError: If upload fails for any reason
130+
"""
131+
try:
132+
# Get URL from environment variable
133+
base_url: str = get_s3_bucket_uri()
134+
135+
filename: str = generate_file_name(img)
136+
137+
full_url = os.path.join(base_url, folder_name, filename)
138+
139+
if img.format is None:
140+
img.format = 'PNG'
141+
142+
mime_type = FORMAT_TO_MIME[img.format.upper()]
143+
144+
buffer: BytesIO = BytesIO()
145+
img_format: str = img.format if img.format else 'PNG'
146+
img.save(buffer, format=img_format)
147+
buffer.seek(0)
148+
149+
aws_request = get_aws_signed_request(full_url, buffer, mime_type).prepare()
150+
151+
response: requests.Response = requests.request(
152+
method=aws_request.method,
153+
url=aws_request.url,
154+
data=aws_request.body,
155+
headers=aws_request.headers,
156+
timeout=30
157+
)
158+
159+
if response.status_code != 200:
160+
raise ImageUploadError(
161+
f"Upload failed with status code {response.status_code}: {response.text}"
162+
)
163+
164+
return full_url
165+
166+
except (InvalidMimeTypeError, MissingEnvironmentVariableError):
167+
raise
168+
except requests.exceptions.RequestException as e:
169+
raise ImageUploadError(f"Network error: {str(e)}")
170+
except Exception as e:
171+
raise ImageUploadError(f"Unexpected error: {str(e)}")
172+
173+
if __name__ == "__main__":
174+
img = Image.new('RGB', (100, 100), color='red')
175+
img.format = 'JPEG'
176+
177+
# Execute
178+
result = upload_image(img, "eduvision")
179+
print(result)

0 commit comments

Comments
 (0)