schort is a tiny link shortener written in python3 and flask
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

schort.py 3.3KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091
  1. #!/usr/bin/env python3
  2. from flask import Flask, render_template, url_for, request, redirect, abort, escape
  3. import sqlite3, random, string, time, hashlib, base64
  4. from urllib.parse import urlparse
  5. app = Flask(__name__)
  6. @app.route('/', methods=['GET', 'POST'])
  7. @app.route('/<shortLink>', methods=['GET', 'POST'])
  8. def short(shortLink=""):
  9. if request.method == "GET":
  10. if shortLink:
  11. noauto = shortLink[-1] == "+"
  12. if noauto: shortLink = shortLink[:-1]
  13. conn = sqlite3.connect("data/links.sqlite")
  14. c = conn.cursor()
  15. result = c.execute('SELECT longLink FROM links WHERE shortLink=?', (shortLink, )).fetchone()
  16. conn.close()
  17. if result:
  18. url = result[0]
  19. parsedUrl = urlparse(url)
  20. if parsedUrl.scheme == "":
  21. url = "http://" + url
  22. if "resolve" in request.args:
  23. return escape(url)
  24. else:
  25. if noauto:
  26. url = str(escape(url))
  27. html = "<a href=" + url + ">" + url + "</a>"
  28. return html
  29. else:
  30. return redirect(url, code=301) # Redirect to long URL saved in the database
  31. else:
  32. return render_template("index.html", name=shortLink, message="Enter long URL for "+ request.url_root + shortLink+":", message_type="info") # Custom link page
  33. else:
  34. return render_template("index.html", name=shortLink) # Landing page
  35. elif request.method == "POST": # Someone submitted a new link to short
  36. longUrl = request.form.get("url", "")
  37. wishId = request.form.get("wishId")
  38. if len(longUrl) <= 0:
  39. abort(400)
  40. databaseId = insertIdUnique(longUrl, idToCheck=wishId)
  41. return request.url_root + databaseId # Short link in plain text
  42. def insertIdUnique(longUrl, idToCheck=None):
  43. hashUrl = hashlib.sha256(longUrl.encode()).digest()
  44. base64Url = base64.urlsafe_b64encode(hashUrl).decode()
  45. if idToCheck == None or idToCheck == "":
  46. idToCheck = base64Url[:4]
  47. conn = sqlite3.connect("data/links.sqlite")
  48. c = conn.cursor()
  49. try:
  50. c.execute('INSERT INTO links VALUES (?, ?, ?, ?, ?)', (idToCheck, longUrl, int(time.time()), request.remote_addr, "default" ))
  51. databaseId = idToCheck
  52. conn.commit()
  53. conn.close()
  54. except sqlite3.IntegrityError as e:
  55. print("Hash already exists, does the long URL matches?")
  56. longUrlDb = c.execute('SELECT * FROM links WHERE shortLink=?', (idToCheck, )).fetchone()
  57. if longUrl == longUrlDb[1]:
  58. print(longUrl + " is already in database with id " + idToCheck + ". Serving old id…")
  59. databaseId = idToCheck
  60. else:
  61. print("Found real hash collision for " + longUrl + " and " + longUrlDb[1])
  62. conn.commit()
  63. conn.close()
  64. if len(base64Url) - 1 >= len(idToCheck) + 1:
  65. databaseId = insertIdUnique(longUrl, idToCheck=base64Url[:len(idToCheck)+1])
  66. else:
  67. print("Can't produce a long enough hash from the new link to be unique. This should never happen")
  68. print("Bailing out, you are on your own. Good luck.")
  69. print("=========================================================================================")
  70. abort(500)
  71. return databaseId
  72. def initDB():
  73. conn = sqlite3.connect("data/links.sqlite")
  74. c = conn.cursor()
  75. c.execute('''CREATE TABLE IF NOT EXISTS links (shortLink UNIQUE NOT NULL, longLink, timestamp, ip, redirectMethod);''')
  76. conn.commit()
  77. conn.close()
  78. if __name__ == '__main__':
  79. initDB()
  80. app.run(debug=True) # If you call this file directly it will always run in debug mode. THIS IS VERY DANGEROUS!
  81. # vim: noexpandtab:ts=2:sw=2:sts=2