From 23a162b71d4eaced5e96548bda446f28ec4e3b4a Mon Sep 17 00:00:00 2001
From: sqozz <sqozz@geekify.de>
Date: Wed, 10 Feb 2016 11:40:36 +0100
Subject: [PATCH] Implemented code that generates IDs out of long urls

---
 shortener.py | 49 +++++++++++++++++++++++++++++++++----------------
 1 file changed, 33 insertions(+), 16 deletions(-)

diff --git a/shortener.py b/shortener.py
index b39891c..0679c7e 100755
--- a/shortener.py
+++ b/shortener.py
@@ -1,12 +1,8 @@
 #!/usr/bin/env python3
-from flask import Flask, render_template, url_for, request, redirect
+from flask import Flask, render_template, url_for, request, redirect, abort
 import sqlite3, random, string, time, hashlib, base64
 app = Flask(__name__)
 
-#@app.route('/')
-#def root():
-#	return render_template("index.html", name=shortLink, message="Enter long URL for "+ request.url_root + shortLink+":", message_type="info")
-
 @app.route('/', methods=['GET', 'POST'])
 @app.route('/<shortLink>', methods=['GET', 'POST'])
 def short(shortLink=""):
@@ -19,25 +15,46 @@ def short(shortLink=""):
 		else:
 				return render_template("index.html", name=shortLink, message="Enter long URL for "+ request.url_root + shortLink+":", message_type="info") # Does the user wish to create a personel short link?
 	elif request.method == "POST": # Someone submitted a new link to short
-		conn = sqlite3.connect("links.sqlite")
-		c = conn.cursor()
 		wishId = request.form["wishId"]
 		longUrl = request.form["url"]
 		if not wishId:
-			hashUrl = hashlib.sha256(longUrl.encode()).hexdigest()
-			base64Url = base64.b64encode(hashUrl.encode()).decode()
-			databaseId = base64Url[:4]
+			databaseId = insertIdUnique("", longUrl)
 		else:
-			databaseId = wishId
+			databaseId = insertIdUnique(wishId, longUrl)
+		return request.url_root + databaseId # TODO: Give the user a nice site where he can see his short URL
 
-		c.execute('INSERT INTO links VALUES (?, ?, ?, ?, ?)', (databaseId, longUrl, int(time.time()), request.remote_addr, "default" ))
+def insertIdUnique(idToCheck, longUrl):
+	hashUrl = hashlib.sha256(longUrl.encode()).hexdigest()
+	base64Url = base64.b64encode(hashUrl.encode()).decode()
+	if len(idToCheck) == 0:
+		idToCheck = base64Url[:4]
+
+	conn = sqlite3.connect("links.sqlite")
+	c = conn.cursor()
+	try:
+		c.execute('INSERT INTO links VALUES (?, ?, ?, ?, ?)', (idToCheck, longUrl, int(time.time()), request.remote_addr, "default" ))
+		databaseId = idToCheck
 		conn.commit()
 		conn.close()
-		return redirect(longUrl, code=302) # TODO: Give the user a nice site where he can see his short URL
+	except sqlite3.IntegrityError as e:
+		print("Hash already exists, does the long URL matches?")
+		longUrlDb = c.execute('SELECT * FROM links WHERE shortLink=?', (idToCheck, )).fetchone()
+		if longUrl == longUrlDb[1]:
+			print(longUrl + " is already in database with id " + idToCheck + ". Serving old id…")
+			databaseId = idToCheck
+		else:
+			print("Found real hash collision for " + longUrl + " and " + longUrlDb[1])
+			conn.commit()
+			conn.close()
+			if len(base64Url) - 1 >= len(idToCheck) + 1:
+				databaseId = insertIdUnique(base64Url[:len(idToCheck)+1], longUrl)
+			else:
+				print("Can't produce a long enough hash from the new link to be unique. This should never happen")
+				print("Bailing out, you are on your own. Good luck.")
+				print("=========================================================================================")
+				abort(500)
 
-def generateId():
-
-	return ''.join(random.SystemRandom().choice(string.ascii_lowercase + string.digits) for _ in range(4))
+	return databaseId
 
 def initDB():
 	conn = sqlite3.connect("links.sqlite")