From 543686dbe8b8fff1cd5f3d53e12c0fc3e76a492a Mon Sep 17 00:00:00 2001
From: xemeds <xemeds@protonmail.com>
Date: Mon, 17 Aug 2020 12:37:00 +0300
Subject: [PATCH] Made new urls not generate a new token if they already exist
 in the database

---
 tiny0/models.py |  2 +-
 tiny0/routes.py | 39 +++++++++++++++++++++++++++++++--------
 tiny0/token.py  |  3 ++-
 3 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/tiny0/models.py b/tiny0/models.py
index c52436f..4326d1c 100644
--- a/tiny0/models.py
+++ b/tiny0/models.py
@@ -3,7 +3,7 @@ from tiny0 import db
 class URL(db.Model):
 	id = db.Column(db.Integer, primary_key=True)
 	token = db.Column(db.String(6), index=True, unique=True, nullable=False)
-	url = db.Column(db.String(2000), nullable=False)
+	url = db.Column(db.String(2000), index=True, unique=True, nullable=False)
 
 	def __repr__(self):
 		return f"'{self.id}' '{self.token}' '{self.url}'"
diff --git a/tiny0/routes.py b/tiny0/routes.py
index 700de53..b8c9bf5 100644
--- a/tiny0/routes.py
+++ b/tiny0/routes.py
@@ -13,17 +13,40 @@ def index():
 
 	# If the form was valid
 	if form.validate_on_submit():
-		# Generate a valid token
-		token = gen_valid_token()
 
-		# Add the token and the given url to the database
-		db.session.add(URL(token=token, url=form.url.data))
-		db.session.commit()
+		# If the given url is a rick roll
+		if ("youtube.com/watch?v=dQw4w9WgXcQ" in form.url.data) or ("youtu.be/dQw4w9WgXcQ" in form.url.data):
+			# Generate a valid token
+			token = gen_valid_token()
 
-		# Return the url page with the shortened url
-		return render_template("url.html", url=WEBSITE_DOMAIN + "/" + token)
+			# Add the token and the given url to the database
+			db.session.add(URL(token=token, url=form.url.data))
+			db.session.commit()
 
-	# If the form was invalid or not submitted
+			# Return the url page with the shortened url
+			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + token)
+
+		# Query the urls that are not a rick roll in the database
+		query = URL.query.filter_by(url=form.url.data).first()
+
+		# If the url exists in the database
+		if query:
+			# Return the url page with the previously shortened url
+			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + query.token)
+
+		# Else if the url does not exist in the database
+		else:
+			# Generate a valid token
+			token = gen_valid_token()
+
+			# Add the token and the given url to the database
+			db.session.add(URL(token=token, url=form.url.data))
+			db.session.commit()
+
+			# Return the url page with the shortened url
+			return render_template("url.html", url=WEBSITE_DOMAIN + "/" + token)
+
+	# Else if the form was invalid or not submitted
 	else:
 		# Return the index page with the form
 		return render_template("index.html", form=form)
diff --git a/tiny0/token.py b/tiny0/token.py
index a592533..00c7adf 100644
--- a/tiny0/token.py
+++ b/tiny0/token.py
@@ -1,3 +1,4 @@
+from tiny0 import db
 from tiny0.models import URL
 from secrets import choice
 
@@ -10,7 +11,7 @@ def gen_valid_token():
 		token = "".join(choice(token_characters) for i in range(6))
 
 		# If the token does not exists in the database
-		if not URL.query.filter_by(token=token).first():
+		if not db.session.query(db.session.query(URL).filter_by(token=token).exists()).scalar():
 			# Break the loop
 			break