2016-02-10 07:57:22 +01:00
#!/usr/bin/env python3
2016-02-10 11:40:36 +01:00
from flask import Flask , render_template , url_for , request , redirect , abort
2016-02-10 07:57:22 +01:00
import sqlite3 , random , string , time , hashlib , base64
app = Flask ( __name__ )
@app.route ( ' / ' , methods = [ ' GET ' , ' POST ' ] )
@app.route ( ' /<shortLink> ' , methods = [ ' GET ' , ' POST ' ] )
def short ( shortLink = " " ) :
if request . method == " GET " :
conn = sqlite3 . connect ( " links.sqlite " )
c = conn . cursor ( )
result = c . execute ( ' SELECT * FROM links WHERE shortLink=? ' , ( shortLink , ) ) . fetchone ( )
if result :
return redirect ( result [ 1 ] , code = 302 ) # Redirect to long URL saved in the database
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
wishId = request . form [ " wishId " ]
longUrl = request . form [ " url " ]
if not wishId :
2016-02-10 11:40:36 +01:00
databaseId = insertIdUnique ( " " , longUrl )
2016-02-10 07:57:22 +01:00
else :
2016-02-10 11:40:36 +01:00
databaseId = insertIdUnique ( wishId , longUrl )
return request . url_root + databaseId # TODO: Give the user a nice site where he can see his short URL
def insertIdUnique ( idToCheck , longUrl ) :
hashUrl = hashlib . sha256 ( longUrl . encode ( ) ) . hexdigest ( )
base64Url = base64 . b64encode ( hashUrl . encode ( ) ) . decode ( )
if len ( idToCheck ) == 0 :
idToCheck = base64Url [ : 4 ]
2016-02-10 07:57:22 +01:00
2016-02-10 11:40:36 +01:00
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
2016-02-10 07:57:22 +01:00
conn . commit ( )
conn . close ( )
2016-02-10 11:40:36 +01:00
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 )
2016-02-10 07:57:22 +01:00
2016-02-10 11:40:36 +01:00
return databaseId
2016-02-10 07:57:22 +01:00
def initDB ( ) :
conn = sqlite3 . connect ( " links.sqlite " )
c = conn . cursor ( )
c . execute ( ''' CREATE TABLE IF NOT EXISTS links (shortLink UNIQUE NOT NULL, longLink, timestamp, ip, redirectMethod); ''' )
conn . commit ( )
conn . close ( )
print ( " DB init " )
if __name__ == ' __main__ ' :
initDB ( )
app . run ( debug = True )