1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
|
from flask_login import (
current_user,
login_user,
login_required,
logout_user,
LoginManager,
UserMixin,
)
from flask import Flask, redirect, render_template, request, url_for, Request
import bcrypt
import os
import string
import random
from models import User
from database import *
from app.util.log import log
class FlaskUser(UserMixin):
pass
app = Flask(__name__)
app.config["SECRET_KEY"] = os.urandom(24)
login_manager = LoginManager()
login_manager.init_app(app)
@login_manager.user_loader
def user_loader(username):
user = FlaskUser()
user.id = username
return user
"""
Handle login requests from the web UI
"""
@app.route("/login", methods=["GET", "POST"])
def login():
if request.method == "POST":
username = request.form["username"]
password = request.form["password"]
# Get database session
db = SessionLocal()
user = db.query(User).filter(User.username == username).first()
db.close()
if not user:
return {"status": "Invalid username or password"}
if bcrypt.checkpw(password.encode("utf-8"), user.password.encode("utf-8")):
flask_user = FlaskUser()
flask_user.id = username
login_user(flask_user)
return {"status": "success"}
return {"status": "Invalid username or password"}
return render_template("login.html")
"""
Handle signup requests from the web UI
"""
@app.route("/signup", methods=["GET", "POST"])
def signup():
if request.method == "POST":
username = request.form["username"]
password = request.form["password"]
# Get database session
db = SessionLocal()
user = db.query(User).filter(User.username == username).first()
if user:
db.close()
return {"status": "User already exists"}
# Add information to the database
hashed_password = bcrypt.hashpw(
password.encode("utf-8"), bcrypt.gensalt()
).decode("utf-8")
api_key = "".join(random.choices(string.ascii_letters + string.digits, k=20))
new_user = User(username=username, password=hashed_password, api_key=api_key)
db.add(new_user)
db.commit()
db.close()
# Log in the newly created user
flask_user = FlaskUser()
flask_user.id = username
login_user(flask_user)
return {"status": "success"}
return render_template("signup.html")
"""
Load the 'dashboard' page for logged in users
"""
@app.route("/dashboard", methods=["GET"])
@login_required
def dashboard():
# Get database session
db = SessionLocal()
# Get the API key for the current user
user = db.query(User).filter(User.username == current_user.id).first()
db.close()
api_key = user.api_key
return render_template("dashboard.html", api_key=api_key)
"""
Log users out of their account
"""
@app.route("/logout", methods=["GET"])
@login_required
def logout():
logout_user()
return redirect(url_for("login"))
"""
Log all records for visits to shortened links
"""
@app.route("/<link>", methods=["GET"])
def log_redirect(link):
# If the `link` is more than 5 characters, ignore
if len(link) > 5:
return
# If the `link` is one of the registered routes, ignore
if link in ["login", "signup", "dashboard", "logout", "api"]:
return
ip = request.remote_addr
user_agent = request.headers.get("user-agent")
redirect_link = log(link, ip, user_agent)
return redirect(redirect_link)
@app.errorhandler(401)
def unauthorized(e):
return redirect(url_for("login"))
@app.errorhandler(404)
def not_found(e):
return redirect(url_for("login"))
|
