From 343fe4357cac314ed7625a1a2b92fffc04f4971d Mon Sep 17 00:00:00 2001
From: Kevin Trocolli <pitok236@gmail.com>
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'}