Browse Source

Implemented code that generates IDs out of long urls

sqozz 3 years ago
parent
commit
23a162b71d
1 changed files with 33 additions and 16 deletions
  1. 33
    16
      shortener.py

+ 33
- 16
shortener.py View File

@@ -1,12 +1,8 @@
1 1
 #!/usr/bin/env python3
2
-from flask import Flask, render_template, url_for, request, redirect
2
+from flask import Flask, render_template, url_for, request, redirect, abort
3 3
 import sqlite3, random, string, time, hashlib, base64
4 4
 app = Flask(__name__)
5 5
 
6
-#@app.route('/')
7
-#def root():
8
-#	return render_template("index.html", name=shortLink, message="Enter long URL for "+ request.url_root + shortLink+":", message_type="info")
9
-
10 6
 @app.route('/', methods=['GET', 'POST'])
11 7
 @app.route('/<shortLink>', methods=['GET', 'POST'])
12 8
 def short(shortLink=""):
@@ -19,25 +15,46 @@ def short(shortLink=""):
19 15
 		else:
20 16
 				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?
21 17
 	elif request.method == "POST": # Someone submitted a new link to short
22
-		conn = sqlite3.connect("links.sqlite")
23
-		c = conn.cursor()
24 18
 		wishId = request.form["wishId"]
25 19
 		longUrl = request.form["url"]
26 20
 		if not wishId:
27
-			hashUrl = hashlib.sha256(longUrl.encode()).hexdigest()
28
-			base64Url = base64.b64encode(hashUrl.encode()).decode()
29
-			databaseId = base64Url[:4]
21
+			databaseId = insertIdUnique("", longUrl)
30 22
 		else:
31
-			databaseId = wishId
23
+			databaseId = insertIdUnique(wishId, longUrl)
24
+		return request.url_root + databaseId # TODO: Give the user a nice site where he can see his short URL
25
+
26
+def insertIdUnique(idToCheck, longUrl):
27
+	hashUrl = hashlib.sha256(longUrl.encode()).hexdigest()
28
+	base64Url = base64.b64encode(hashUrl.encode()).decode()
29
+	if len(idToCheck) == 0:
30
+		idToCheck = base64Url[:4]
32 31
 
33
-		c.execute('INSERT INTO links VALUES (?, ?, ?, ?, ?)', (databaseId, longUrl, int(time.time()), request.remote_addr, "default" ))
32
+	conn = sqlite3.connect("links.sqlite")
33
+	c = conn.cursor()
34
+	try:
35
+		c.execute('INSERT INTO links VALUES (?, ?, ?, ?, ?)', (idToCheck, longUrl, int(time.time()), request.remote_addr, "default" ))
36
+		databaseId = idToCheck
34 37
 		conn.commit()
35 38
 		conn.close()
36
-		return redirect(longUrl, code=302) # TODO: Give the user a nice site where he can see his short URL
37
-
38
-def generateId():
39
+	except sqlite3.IntegrityError as e:
40
+		print("Hash already exists, does the long URL matches?")
41
+		longUrlDb = c.execute('SELECT * FROM links WHERE shortLink=?', (idToCheck, )).fetchone()
42
+		if longUrl == longUrlDb[1]:
43
+			print(longUrl + " is already in database with id " + idToCheck + ". Serving old id…")
44
+			databaseId = idToCheck
45
+		else:
46
+			print("Found real hash collision for " + longUrl + " and " + longUrlDb[1])
47
+			conn.commit()
48
+			conn.close()
49
+			if len(base64Url) - 1 >= len(idToCheck) + 1:
50
+				databaseId = insertIdUnique(base64Url[:len(idToCheck)+1], longUrl)
51
+			else:
52
+				print("Can't produce a long enough hash from the new link to be unique. This should never happen")
53
+				print("Bailing out, you are on your own. Good luck.")
54
+				print("=========================================================================================")
55
+				abort(500)
39 56
 
40
-	return ''.join(random.SystemRandom().choice(string.ascii_lowercase + string.digits) for _ in range(4))
57
+	return databaseId
41 58
 
42 59
 def initDB():
43 60
 	conn = sqlite3.connect("links.sqlite")

Loading…
Cancel
Save