Implemented code that generates IDs out of long urls

This commit is contained in:
sqozz 2016-02-10 11:40:36 +01:00
parent 551d9e1cce
commit 23a162b71d

View file

@ -1,12 +1,8 @@
#!/usr/bin/env python3 #!/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 import sqlite3, random, string, time, hashlib, base64
app = Flask(__name__) 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('/', methods=['GET', 'POST'])
@app.route('/<shortLink>', methods=['GET', 'POST']) @app.route('/<shortLink>', methods=['GET', 'POST'])
def short(shortLink=""): def short(shortLink=""):
@ -19,25 +15,46 @@ def short(shortLink=""):
else: 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? 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 elif request.method == "POST": # Someone submitted a new link to short
conn = sqlite3.connect("links.sqlite")
c = conn.cursor()
wishId = request.form["wishId"] wishId = request.form["wishId"]
longUrl = request.form["url"] longUrl = request.form["url"]
if not wishId: if not wishId:
hashUrl = hashlib.sha256(longUrl.encode()).hexdigest() databaseId = insertIdUnique("", longUrl)
base64Url = base64.b64encode(hashUrl.encode()).decode()
databaseId = base64Url[:4]
else: 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.commit()
conn.close() 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 databaseId
return ''.join(random.SystemRandom().choice(string.ascii_lowercase + string.digits) for _ in range(4))
def initDB(): def initDB():
conn = sqlite3.connect("links.sqlite") conn = sqlite3.connect("links.sqlite")