From 343fe4357cac314ed7625a1a2b92fffc04f4971d Mon Sep 17 00:00:00 2001 From: Kevin Trocolli Date: Sun, 16 Jul 2023 16:58:18 -0400 Subject: [PATCH] mai2: add image validation via Pillow --- requirements.txt | 1 + titles/mai2/base.py | 46 +++++++++++++++++++++++++++++++++++---------- 2 files changed, 37 insertions(+), 10 deletions(-) diff --git a/requirements.txt b/requirements.txt index 6d37728..53d867a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -17,3 +17,4 @@ bcrypt jinja2 protobuf autobahn +pillow diff --git a/titles/mai2/base.py b/titles/mai2/base.py index 267edec..f1b04b6 100644 --- a/titles/mai2/base.py +++ b/titles/mai2/base.py @@ -2,7 +2,8 @@ from datetime import datetime from typing import Any, Dict, List import logging from base64 import b64decode -from os import path, stat +from os import path, stat, remove +from PIL import ImageFile from core.config import CoreConfig from titles.mai2.const import Mai2Constants @@ -91,7 +92,7 @@ class Mai2Base: for i, charge in enumerate(game_charge_list): charge_list.append( { - "orderId": i, + "orderId": i + 1, "chargeId": charge["ticketId"], "price": charge["price"], "startDate": "2017-12-05 07:00:00.0", @@ -812,27 +813,52 @@ class Mai2Base: self.logger.warn(f"Photo too large ({div_len} * 10240 = {div_len * 10240} bytes)") return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} + ret_code = order_id + 1 photo_chunk = b64decode(div_data) if len(photo_chunk) > 10240 or (len(photo_chunk) < 10240 and div_num + 1 != div_len): self.logger.warn(f"Incorrect data size after decoding (Expected 10240, got {len(photo_chunk)})") return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} + + out_name = f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}" - if not path.exists(f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}.jpeg") and div_num != 0: + if not path.exists(f"{out_name}.bin") and div_num != 0: self.logger.warn(f"Out of order photo upload (div_num {div_num})") return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} - if path.exists(f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}.jpeg") and div_num == 0: + if path.exists(f"{out_name}.bin") and div_num == 0: self.logger.warn(f"Duplicate file upload") return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} - elif path.exists(f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}.jpeg"): - fstats = stat(f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}.jpeg") + elif path.exists(f"{out_name}.bin"): + fstats = stat(f"{out_name}.bin") if fstats.st_size != 10240 * div_num: self.logger.warn(f"Out of order photo upload (trying to upload div {div_num}, expected div {fstats.st_size / 10240} for file sized {fstats.st_size} bytes)") return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} - - with open(f"{self.game_config.uploads.photos_dir}/{user_id}_{playlog_id}_{track_num}.jpeg", "ab") as f: - f.write(photo_chunk) + + try: + with open(f"{out_name}.bin", "ab") as f: + f.write(photo_chunk) + + except Exception: + self.logger.error(f"Failed writing to {out_name}.bin") + return {'returnCode': 0, 'apiName': 'UploadUserPhotoApi'} - return {'returnCode': order_id + 1, 'apiName': 'UploadUserPhotoApi'} + if div_num + 1 == div_len and path.exists(f"{out_name}.bin"): + try: + p = ImageFile.Parser() + with open(f"{out_name}.bin", "rb") as f: + p.feed(f.read()) + + im = p.close() + im.save(f"{out_name}.jpeg") + except Exception: + self.logger.error(f"File {out_name}.bin failed image validation") + + try: + remove(f"{out_name}.bin") + + except Exception: + self.logger.error(f"Failed to delete {out_name}.bin, please remove it manually") + + return {'returnCode': ret_code, 'apiName': 'UploadUserPhotoApi'}