From 25fb22bc5c1c91e02d37b276a2eb463964656aa8 Mon Sep 17 00:00:00 2001
From: James Wang <jameswang9909@hotmail.com>
Date: Sat, 2 Jan 2016 15:45:05 -0500
Subject: [PATCH] Add files model to database

---
 .gitignore            |  1 +
 server/api/models.py  |  9 +++++++++
 server/api/problem.py | 23 +++++++++++++++++++++--
 server/app.py         |  3 ++-
 server/config.py      |  2 ++
 5 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/.gitignore b/.gitignore
index 66f8da0..bb191d4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@
 .secret_key
 .bundle/config
 logs/
+uploads/
 
 # Object files
 *.o
diff --git a/server/api/models.py b/server/api/models.py
index d141404..e80064e 100644
--- a/server/api/models.py
+++ b/server/api/models.py
@@ -53,6 +53,15 @@ class Problems(db.Model):
         self.flag = flag
         self.value = value
 
+class Files(db.Model):
+    fid = db.Column(db.Integer, primary_key=True)
+    pid = db.Column(db.Integer)
+    location = db.Column(db.Text)
+
+    def __init__(self, pid, location):
+        self.pid = pid
+        self.location = location
+
 class Solves(db.Model):
     sid = db.Column(db.Integer, primary_key=True)
     pid = db.Column(db.Integer)
diff --git a/server/api/problem.py b/server/api/problem.py
index d55a1f5..5b8d857 100644
--- a/server/api/problem.py
+++ b/server/api/problem.py
@@ -1,9 +1,11 @@
+import hashlib
 import logger
 
 from flask import Blueprint, session, request
 from flask import current_app as app
+from werkzeug import secure_filename
 
-from models import db, Problems, Solves, Teams
+from models import db, Files, Problems, Solves, Teams
 from decorators import admins_only, api_wrapper, login_required
 
 blueprint = Blueprint("problem", __name__)
@@ -20,12 +22,29 @@ def problem_add():
     value = request.form["value"]
 
     name_exists = Problems.query.filter_by(name=name).first()
-
     if name_exists:
         return { "success":0, "message": "Problem name already taken." }
 
     problem = Problems(name, category, description, hint, flag, value)
     db.session.add(problem)
+
+    files = request.files["imp-files"]
+    for _file in files:
+        filename = secure_filename(_file.filename)
+
+        if len(filename) == 0:
+            continue
+
+        folder = problem.name.replace(" ", "-")
+        folder_path = os.path.join(os.path.normpath(app["UPLOAD_FOLDER"], folder))
+        if not folder_path:
+            os.makedirs(folder_path)
+
+        file_path = os.path.join(path, filename)
+        _file.save(file_path)
+        db_file = Files(problem.pid, file_path)
+        db.session.add(db_file)
+
     db.session.commit()
 
     return { "success": 1, "message": "Success!" }
diff --git a/server/app.py b/server/app.py
index e5a2378..7ba7b69 100644
--- a/server/app.py
+++ b/server/app.py
@@ -9,9 +9,10 @@ app = Flask(__name__)
 
 app.config["SQLALCHEMY_DATABASE_URI"] = config.SQLALCHEMY_DATABASE_URI
 app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = config.SQLALCHEMY_TRACK_MODIFICATIONS
+app.config["UPLOAD_FOLDER"] = config.UPLOAD_FOLDER
 
 with app.app_context():
-	from api.models import db, Teams, Problems, Solves, Users
+	from api.models import db, Files, Teams, Problems, Solves, Users
 	db.init_app(app)
 	db.create_all()
 
diff --git a/server/config.py b/server/config.py
index a24d426..8d3112c 100644
--- a/server/config.py
+++ b/server/config.py
@@ -15,5 +15,7 @@ SECRET_KEY = key
 SQLALCHEMY_DATABASE_URI = "mysql://root:i_hate_passwords@localhost/easyctf"
 SQLALCHEMY_TRACK_MODIFICATIONS = False
 
+UPLOAD_FOLDER = os.path.normpath("uploads")
+
 CTF_BEGIN = 0 # To be used later
 CTF_END = 0 # To be used later