From 9abb8aa47fd864f2d1ed65deb538c44de1174e9c Mon Sep 17 00:00:00 2001
From: Carlos Quintana <74399022+cquintana92@users.noreply.github.com>
Date: Wed, 29 Jun 2022 15:04:55 +0200
Subject: [PATCH] Validate user uploaded image (#1123)

* Validate user uploaded image

* Fix test/data path detection
---
 app/dashboard/views/setting.py |  13 ++++++---
 app/image_validation.py        |  28 ++++++++++++++++++++
 tests/data/1px.jpg             | Bin 0 -> 631 bytes
 tests/data/1px.webp            | Bin 0 -> 44 bytes
 tests/test_image_validation.py |  47 +++++++++++++++++++++++++++++++++
 5 files changed, 85 insertions(+), 3 deletions(-)
 create mode 100644 app/image_validation.py
 create mode 100644 tests/data/1px.jpg
 create mode 100644 tests/data/1px.webp
 create mode 100644 tests/test_image_validation.py

diff --git a/app/dashboard/views/setting.py b/app/dashboard/views/setting.py
index c08ab145..28ea408d 100644
--- a/app/dashboard/views/setting.py
+++ b/app/dashboard/views/setting.py
@@ -29,6 +29,7 @@ from app.email_utils import (
     personal_email_already_used,
 )
 from app.errors import ProtonPartnerNotSetUp
+from app.image_validation import detect_image_format, ImageFormat
 from app.jobs.export_user_data_job import ExportUserDataJob
 from app.log import LOG
 from app.models import (
@@ -181,12 +182,18 @@ def setting():
                     profile_updated = True
 
                 if form.profile_picture.data:
+                    image_contents = form.profile_picture.data.read()
+                    if detect_image_format(image_contents) == ImageFormat.Unknown:
+                        flash(
+                            "This image format is not supported",
+                            "error",
+                        )
+                        return redirect(url_for("dashboard.setting"))
+
                     file_path = random_string(30)
                     file = File.create(user_id=current_user.id, path=file_path)
 
-                    s3.upload_from_bytesio(
-                        file_path, BytesIO(form.profile_picture.data.read())
-                    )
+                    s3.upload_from_bytesio(file_path, BytesIO(image_contents))
 
                     Session.flush()
                     LOG.d("upload file %s to s3", file)
diff --git a/app/image_validation.py b/app/image_validation.py
new file mode 100644
index 00000000..6c5bbb2a
--- /dev/null
+++ b/app/image_validation.py
@@ -0,0 +1,28 @@
+from enum import Enum
+
+
+class ImageFormat(Enum):
+    Png = 1
+    Jpg = 2
+    Webp = 3
+    Svg = 4
+    Unknown = 9
+
+
+magic_numbers = {
+    ImageFormat.Png: bytes([0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A]),
+    ImageFormat.Jpg: bytes([0xFF, 0xD8, 0xFF, 0xE0]),
+    ImageFormat.Webp: bytes([0x52, 0x49, 0x46, 0x46]),
+    ImageFormat.Svg: bytes([0x3C, 0x3F, 0x78, 0x6D, 0x6C]),  # <?xml
+}
+
+
+def detect_image_format(image: bytes) -> ImageFormat:
+    # Detect image based on magic number
+    for fmt, header in magic_numbers.items():
+        if image.startswith(header):
+            return fmt
+    # Detect if is svg
+
+    # We don't know the type
+    return ImageFormat.Unknown
diff --git a/tests/data/1px.jpg b/tests/data/1px.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..1cda9a53dc357ce07d3c67051b7615ebf7dc2f64
GIT binary patch
literal 631
zcmex=<NpH&0WUXCHwH#VMur3+WcYuZ!I^=Bjg6g+m4ls~os*M;i${c)hnt&6Qb?Fz
zL{>^(PF6}rMnOeST|r4lSw=>~TvNxu(8R<<Uft5x!pKI?*u)5A2qPyaCl5D|1TU|I
zk)n*E5y{~H0R}-121W)(W=16jCP7AKLB{__7~~lk7+Ha?1vwZB7@3$^SlQS)IJvli
z3bqO`FflSSGqEtUva+xMRo4RL8CV2ag%k}P*@OcV*_8@Kj2b5{<WP3ncu+Lx;s+Ju
zq@pHHE-`TlNhwt|bq!4|6H_yD3rj0!7gslT4^OY)kkGL3h{&kql+?8JjLfX!lG3vB
zipr|yme#iRj?S)0lc!9bHhsp-S&J4gS-Ncbij}K2ZQinV+x8thcO5!><mj>ECr+Na
zbot8FYu9hwy!G(W<0ns_J%91?)yGetzkL1n{m0K=Ab&A3Fhjfr_ZgbM1cClyVqsxs
zVF&q(k*OSrnFU!`6%E;h90S=C3x$=88aYIqCNA7~kW<+>=!0ld(M2vX6_bamA3<IN
m`;0h`HId~rxW^Fwy2Zf5%m|D;W<dsfhVQ>L7B$%azX<@d&d)*s

literal 0
HcmV?d00001

diff --git a/tests/data/1px.webp b/tests/data/1px.webp
new file mode 100644
index 0000000000000000000000000000000000000000..c78718cda3c6ea710aaf14e083d87b10d62bc6bb
GIT binary patch
literal 44
ycmWIYbaPW-U|<M$bqWXzuuzZyvJDs+<}zwAGB7eQF_@?>VPGy``1gAX0|Nl9^$3ap

literal 0
HcmV?d00001

diff --git a/tests/test_image_validation.py b/tests/test_image_validation.py
new file mode 100644
index 00000000..51e5c78e
--- /dev/null
+++ b/tests/test_image_validation.py
@@ -0,0 +1,47 @@
+from app.image_validation import ImageFormat, detect_image_format
+from pathlib import Path
+
+
+def get_path_to_static_dir() -> Path:
+    this_path = Path(__file__)
+    repo_root_path = this_path.parent.parent
+    return repo_root_path.joinpath("static")
+
+
+def read_static_file_contents(filename: str) -> bytes:
+    image_path = get_path_to_static_dir().joinpath(filename)
+    with open(image_path.as_posix(), "rb") as f:
+        return f.read()
+
+
+def read_test_data_file_contents(filename: str) -> bytes:
+    this_path = Path(__file__)
+    test_data_path = this_path.parent.joinpath("data")
+    file_path = test_data_path.joinpath(filename)
+    with open(file_path.as_posix(), "rb") as f:
+        return f.read()
+
+
+def test_non_image_file_returns_unknown():
+    contents = read_static_file_contents("local-storage-polyfill.js")
+    assert detect_image_format(contents) is ImageFormat.Unknown
+
+
+def test_png_file_is_detected():
+    contents = read_static_file_contents("logo.png")
+    assert detect_image_format(contents) is ImageFormat.Png
+
+
+def test_jpg_file_is_detected():
+    contents = read_test_data_file_contents("1px.jpg")
+    assert detect_image_format(contents) is ImageFormat.Jpg
+
+
+def test_webp_file_is_detected():
+    contents = read_test_data_file_contents("1px.webp")
+    assert detect_image_format(contents) is ImageFormat.Webp
+
+
+def test_svg_file_is_detected():
+    contents = read_static_file_contents("icon.svg")
+    assert detect_image_format(contents) is ImageFormat.Svg