From 699fe59388b2068a89a183a35594858df4a219bb Mon Sep 17 00:00:00 2001
From: Mateusz779 <mateusz.kedziora@mkedziora.pl>
Date: Fri, 7 Apr 2023 13:25:24 +0200
Subject: [PATCH] =?UTF-8?q?podstawowa=20obs=C5=82uga=20bazy=20danych?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 app.py    | 60 ++++++++++++++++++++++++++++++++++++++++++++++++++++---
 config.py |  5 +++++
 db.py     | 50 ++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+), 3 deletions(-)
 create mode 100644 config.py
 create mode 100644 db.py

diff --git a/app.py b/app.py
index c2f2d7f..5383208 100644
--- a/app.py
+++ b/app.py
@@ -1,7 +1,61 @@
-from flask import Flask, send_file
+from flask import Flask, send_file, jsonify, request
+import db
+import os
+from werkzeug.utils import secure_filename
 
 app = Flask(__name__)
+app.config['UPLOAD_FOLDER'] = "configs/squash"
+app.config['MAX_CONTENT_LENGTH'] = 1024 * 1024 * 512 #512MB
+
+@app.route("api/addimage", methods=['POST'])
+def add_image():
+    db.Connect()
+    name = None
+    try:
+        file = request.files['file']
+        if file is None or file == "":
+            return jsonify(message="nofile")
+    except Exception as e:
+        return jsonify(message="nofile")
+    
+    try:
+        token = request.form['token']
+        if token is None or token == "":
+            return jsonify(message="notoken")
+    except:
+        if token is None:
+            return jsonify(message="notoken")
+    
+    incorrect = True
+    while incorrect:
+        if db.GetVPNImage(token) is not None:
+            if name[-1:].isdigit():
+                name = name[:-1] + str(int(name[-1:])+1)
+            else:
+                name = name+"1"
+        else:
+            incorrect = False
+
+    filename = secure_filename(file.filename)
+    while os.path.exists(os.path.join(app.config['UPLOAD_FOLDER'], filename)):
+        if filename[0].isdigit():
+            filename = str(int(filename[0])+1)+filename[1:]
+        else:
+            filename = "1"+filename
+    file.save(os.path.join(app.config['UPLOAD_FOLDER'], filename))
+    
+    db.AddVPNImage(filename, token)
+
+    return jsonify(message="ok")
+   
 
 @app.route("/api/getvpn")
-def hello_world():
-    return send_file("configs/squash/praktyki.squashfs")
\ No newline at end of file
+def get_image():
+    try:
+        filename = db.GetVPNImage(request.args['token'])
+    except:
+        filename = "default.squashfs"
+    if filename is None or filename == "":
+        filename = "default.squashfs"
+        
+    return send_file(os.path.join(app.config['UPLOAD_FOLDER'], filename))
\ No newline at end of file
diff --git a/config.py b/config.py
new file mode 100644
index 0000000..2038351
--- /dev/null
+++ b/config.py
@@ -0,0 +1,5 @@
+database="praktyki"
+host="localhost"
+user="praktyki"
+password="2a7driUITXFy73tO"
+port="5432"
\ No newline at end of file
diff --git a/db.py b/db.py
new file mode 100644
index 0000000..e260539
--- /dev/null
+++ b/db.py
@@ -0,0 +1,50 @@
+import psycopg2
+import config
+
+def Connect():
+    global cur, conn
+    try:
+        conn = psycopg2.connect(database=config.database,
+                                host=config.host,
+                                user=config.user,
+                                password=config.password,
+                                port=config.port)
+    except Exception as ex:
+        print(f"Error connecting to PostgreSQL: {e}")
+        
+    cur = conn.cursor()
+    
+    with conn.cursor() as cur:
+        cur.execute("""
+            CREATE TABLE IF NOT EXISTS vpn (
+                id SERIAL PRIMARY KEY,
+                image_name VARCHAR(255) NOT NULL,
+                token VARCHAR(255) NOT NULL,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+        """)
+        conn.commit()
+    
+def GetCur():
+    return conn.cursor()
+
+def GetConn():
+    return conn
+
+def AddVPNImage(name, token):
+    Connect()
+    with GetCur() as cur:
+        cur.execute("""
+            INSERT INTO vpn (image_name, token)
+            VALUES (%s, %s)
+        """,(name, token,))
+        conn.commit()
+        
+def GetVPNImage(token):
+    Connect()
+    with GetCur() as cur:
+        return cur.execute("""
+            SELECT image_name FROM vpn WHERE token = %s
+        """,(token,)).fetchone()
+        
+    
\ No newline at end of file