From c52b4178403b641c4383bafc41be44ae4e669394 Mon Sep 17 00:00:00 2001
From: doobry <doobry@systemli.org>
Date: Tue, 26 Feb 2019 12:40:11 +0100
Subject: [PATCH] Add support to override default templates with custom ones
 (Fixes: #14)

---
 .gitignore | 1 +
 README.md  | 6 ++++++
 main.py    | 7 +++++++
 3 files changed, 14 insertions(+)

diff --git a/.gitignore b/.gitignore
index bee8a64..e42d897 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 __pycache__
+custom_templates
diff --git a/README.md b/README.md
index b12e1ea..5f00d86 100644
--- a/README.md
+++ b/README.md
@@ -79,6 +79,12 @@ collector cronjob to remove leftover files . Besides, it can create a
 the uploads folder, to ensure that the uploaded files won't be recoverable
 between reboots.
 
+# Custom templates
+
+You can override the default templates from `templates/` by putting replacements
+into the directory path that's configured in `app.config['CUSTOM_TEMPLATES_DIR']`
+(default `custom_templates/`).
+
 # Threat model
 
 - An attacker in possession of the very same file that a user wants to clean,
diff --git a/main.py b/main.py
index 59c3791..033d8b6 100644
--- a/main.py
+++ b/main.py
@@ -6,6 +6,7 @@ from libmat2 import parser_factory
 
 from flask import Flask, flash, request, redirect, url_for, render_template
 from flask import send_from_directory, after_this_request
+import jinja2
 
 from werkzeug.utils import secure_filename
 
@@ -14,6 +15,12 @@ app = Flask(__name__)
 app.config['SECRET_KEY'] = os.urandom(32)
 app.config['UPLOAD_FOLDER'] = './uploads/'
 app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB
+app.config['CUSTOM_TEMPLATES_DIR'] = 'custom_templates'
+
+app.jinja_loader = jinja2.ChoiceLoader([ # type: ignore
+    jinja2.FileSystemLoader(app.config['CUSTOM_TEMPLATES_DIR']),
+    app.jinja_loader,
+    ])
 
 def __hash_file(filepath: str) -> str:
     sha256 = hashlib.sha256()