Reset Database Function Added!

Author: thevamp99

.idea/TIWAP.iml

@@ -117,8 +117,16 @@ def settings():
 
         print(session['level'])
 
-        return render_template('settings.html', level=level, msg="Difficult Set to " + level)
+        return render_template('settings.html', level=level, msg="Difficulty Set to " + level)
 
+# Reset DB
+@app.route('/reset-db')
+@is_logged
+def reset_db():
+    dbm.reset_db()
+    mongo_dbm.reset_db()
+
+    return render_template('settings.html', msg="Database has been reset!")
 
 # SQL Injection
 @app.route('/sql-injection', methods=['POST', 'GET'])

.idea/modules.xml

@@ -106,3 +106,14 @@ def save_name(self, name):
             return True
 
         return False
+
+    def reset_db(self):
+        self.create_db_connection()
+        with open('helper/sqlite_db_reset.txt') as f:
+            contents = f.readlines()
+
+        for content in contents:
+            self.cur.execute(content)
+
+        self.close_db_connection()
+        self.commit_db()

.idea/vcs.xml

@@ -1,4 +1,5 @@
 from pymongo import MongoClient
+import pymongo
 
 # Global Variables
 client = MongoClient("mongodb://localhost:27017/")
@@ -28,3 +29,20 @@ def get_data_filtered(self, filtered):
 
         except Exception as e:
             return e
+
+    def reset_db(self):
+        db = client["TIWAP"]
+        col_1 = db["cars"]
+        col_2 = db["users"]
+
+        col_1.delete_many({})
+        col_2.delete_many({})
+
+        dict_1 = {"name":"720d", "company":"BMW"}
+        dict_2 = {'name': 'G63','company': 'AMG'}
+        dict_3 = {'name': 'A8','company': 'Audi'}
+        dict_4 = {'username':'admin','password':'admin'}
+        dict_5 = {'username':'john','password':'john123'}
+
+        col_1.insert_many([dict_1,dict_2,dict_3])
+        col_2.insert_many([dict_4,dict_5])

TIWAF.db

@@ -0,0 +1,10 @@
+DROP TABLE comments
+DROP TABLE names
+CREATE TABLE comments ( comment TEXT )
+CREATE TABLE names ( name TEXT )
+INSERT INTO comments VALUES('This Project is so cool.')
+INSERT INTO comments VALUES('Having Fun with XSS.')
+INSERT INTO comments VALUES('This is hacked!')
+INSERT INTO names VALUES('Mark')
+INSERT INTO names VALUES('Jeff')
+INSERT INTO names VALUES('Sundar')

app.py

@@ -74,6 +74,10 @@ <h1 id="vuln-name">Instructions</h1>
             </select>
             <button class="btn btn-success" type="submit">Set Level</button>
         </form>
+        <form action="/reset-db">
+            <p><br><b>Reset the Database</b></p>
+            <button class="btn btn-success" type="submit">Reset DB</button>
+        </form>
         <h6 class="text-center text-muted mt-3">{{ msg }}</h6>
     </div>
 </div>