main.py

#!/usr/bin/env python3
from waitress import serve #Used as webserver (Production)
from flask import Flask, request, render_template, redirect, abort, Markup #Used to prepare the dynamic pages (The main site)
import sqlite3 #Used to store the Data
import os #Used for getting the enviorement variables
import qrcode #Used to generate the QR
import base64 #Used to encode the generated QR as base64, to directly insert it into the HTML
from io import BytesIO #Needed for base64 encoding of the image
from PIL import Image #Needed for QR generation

app = Flask(__name__)
domain_to_index = {}
domain_prepared = ""

try:
    domain = os.environ["domains"].split(";") #Get the domains from the enviorement variable. If no enviorement variable is set set it to 127.0.0.1:5000 (for testing)
except:
    domain = ["127.0.0.1:5000"]

builddate = ""
try:
    if(os.environ["show_build_date"] == "1"): #If you want to see the builddate you can enable this enviorement variable
        builddate = open("builddate.txt", "r").read()
except:
    pass #This exception is only to pass it while testing

try:
    if(os.environ["production"] == "1"): #If you use this in production, please set this to 1, because the Flask Testserver is not very secure (e.g. shows error on Website)
        production = True
    else:
        production = False
except:
    production = False

try:
    url_scheme = os.environ["url_scheme"]
except:
    url_scheme = "http"

index = 0
for domains in domain: #Make from every domnain a entry for the select box later
    domains = domains
    domain_prepared = domain_prepared + '<option value="' + str(domains) + '">' + str(domains) + '</option>'
    domain_to_index[domains] = str(index)
    index = index + 1


def table_check():
    create_table = """
        CREATE TABLE WEB_URL(
        LONG_URL TEXT NOT NULL,
        SHORT_URL TEXT NOT NULL
        );
        """
    with sqlite3.connect('db/urls.db') as conn:
        cursor = conn.cursor()
        try: #Try making the database structure, if fails Database was already created.
            cursor.execute(create_table)
        except sqlite3.OperationalError:
            pass


def makeQR(text): #This function is used to create a QR code and encode it base64, if you make a new shortlink
    qr = qrcode.QRCode( #QR generation variables
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=10,
        border=1,
    )
    qr.add_data(text) #The URL is in the text variable
    qr.make(fit=True) #Generate the QR

    img = qr.make_image(fill_color="black", back_color="white") #Encode the WR as base 64
    with BytesIO() as buffer:
        img.save(buffer, 'jpeg')
        return base64.b64encode(buffer.getvalue()).decode()

@app.route('/', methods=['GET'])
def home_get():
    return render_template('home.html', builddate=builddate, domain=domain_prepared) #return the default site to create a new shorten link

@app.route('/', methods=['POST']) #This function is used to create a new url
def home_post():
    if (request.form.get('url').replace(" ", "") == ""):
        return render_template('home.html', builddate=builddate, domain=domain_prepared, snackbar="Please enter a url to short, before submitting this form", long_url_prefilled=request.form.get('url'), short_url_prefilled=request.form.get('short').lower(), domain_prefilled=domain_to_index[request.form.get('domain')]) #return the user the prefilled form with an error message, because no url to short was provided
    if (request.form.get('short').replace(" ", "") == ""):
        return render_template('home.html', builddate=builddate, domain=domain_prepared, snackbar="Please enter a short name, before submitting this form", long_url_prefilled=request.form.get('url'), short_url_prefilled=request.form.get('short').lower(), domain_prefilled=domain_to_index[request.form.get('domain')]) #return the user the prefilled form with an error message, because no short link was provided
    shorturl = (request.form.get('domain') + "/" + request.form.get('short')).lower()

    url = request.form.get('url')
    with sqlite3.connect('db/urls.db') as conn: #Check if another user already used the short link
        cursor = conn.cursor()
        res = cursor.execute('SELECT LONG_URL FROM WEB_URL WHERE SHORT_URL=?', [shorturl])
        try:
            short = res.fetchone()
            already_used = False
            if short is not None:
                already_used = True
        except:
            pass

        if not already_used: #If short link wasn't used before, insert the link in the Database.
            res = cursor.execute(
                'INSERT INTO WEB_URL (LONG_URL, SHORT_URL) VALUES (?, ?)',
                [url, shorturl]
            )
            return render_template('home.html', short_url=shorturl, builddate=builddate, domain=domain_prepared, qrcode=makeQR("http://" + shorturl)) #return the shorten link to the user
        else:
            return render_template('home.html', builddate=builddate, domain=domain_prepared, snackbar="URL already used, please try another one", long_url_prefilled=request.form.get('url'), short_url_prefilled=request.form.get('short').lower(), domain_prefilled=domain_to_index[request.form.get('domain')]) #return the user the prefilled form with an error message, because the url was already used
    
@app.route('/favicon.ico') #Redirect to the static url of the favicon
def favicon():
    return redirect("/static/favicon.ico")

@app.route('/<short_url>')
def redirect_short_url(short_url):
    host = request.headers['Host']
    url = ""
    with sqlite3.connect('db/urls.db') as conn: #Get the original URL from the database
        cursor = conn.cursor()
        res = cursor.execute('SELECT LONG_URL FROM WEB_URL WHERE SHORT_URL=?', [host + "/" + short_url.lower()])
        try:
            short = res.fetchone()
            if short is not None: #If a long url is found
                url = short[0]
                error_404 = False
            else:
                error_404 = True #If no url is found throw a 404. If you throw a error in a try / catch block it will be catched by this, so set a variable to true and throw the error later
        except Exception as e: #If there happens an error, print the exception to the console and throw a 500 error
            print(e) #Print a debug Message to the console
            abort(500) #Throw a 500 error. This means internal Server error.
    if not error_404: #If there was no 404 error before, redirect the user. If not throw a 404 error
        return redirect(url)
    else:
        abort(404)


if __name__ == '__main__':
    table_check()# This code checks whether database table is created or not
    if production: #Check if production variable is set to true use the waitress webserver, else use the buildin flask webserver, with more debug output
        serve(app, host='0.0.0.0', port= 5000, url_scheme=url_scheme) #Start the Webserver for all users on port 5000
    else:
        app.run(host='0.0.0.0', port=5000, debug=True) #Start the Webserver in Debug mode. This means, if the script runs in an error, it will show the error message in Browser.