from app import app, worker, bcrypt, loginManager from app.models import User from config import Config, Errors from flask import render_template, request, send_file, redirect, flash from flask_login import login_user, current_user, logout_user, login_required from flask_wtf import FlaskForm from wtforms import StringField, PasswordField, SubmitField, BooleanField from wtforms.validators import DataRequired, Length, EqualTo from werkzeug.datastructures import FileStorage from werkzeug.utils import secure_filename from io import BytesIO import os import io import random import magic class RegistrationForm(FlaskForm): username = StringField('Username', validators=[DataRequired(), Length(min=2, max=16)]) password = PasswordField('Password', validators=[DataRequired(), Length(min=8, max=32)]) password2 = PasswordField('Confirm Password', validators=[DataRequired(), EqualTo('password')]) tnc = BooleanField('I agree to the Terms and Conditions', validators=[DataRequired()]) submit = SubmitField('Register') def validate_username(self, username): user = Config.users.find_one({"username": username.data}) if user: raise ValueError("That username is taken. Try another.") class LoginForm(FlaskForm): username = StringField('Username', validators=[DataRequired(), Length(min=2, max=16)]) password = PasswordField('Password', validators=[DataRequired(), Length(min=8, max=32)]) submit = SubmitField('Login') @loginManager.user_loader def load_user(userid): user = User.get(userid) return user @app.route('/', methods=["GET", "POST"]) def index(): # Check for a GET or POST request if request.method == "GET": print(current_user.is_authenticated) return render_template('index.html') elif request.method == "POST": # Before anything else, we want to take the IP if the logging is enabled if Config.ipLogEnabled == True: ip = request.remote_addr else: # If not then return a 0 ip = 0 # Init variables before they're passed userid = request.form.get("userid") if request.form.get("userid") else None filename = request.form.get("filename") if request.form.get("filename") else None retention = request.form.get("retention") if request.form.get("retention") else None id = request.form.get("filename") if Config.files.find_one({"id": filename}) is None else None # We got a file or a url? if 'file' in request.files: # Grab the file and store it, this is a FileStorage object file = request.files['file'] # Call the function to upload the file, this will return either HTTP Status codes or a 200 with a URL. result, status = worker.uploadFile(file, ip, userid, filename, id, retention) result = "https://xygt.cc/{}".format(result) return result, status elif 'file' in request.form: file = FileStorage(stream=BytesIO(request.form['file'].encode("utf-8")), filename=id, content_type="text/plain") result, status = worker.uploadFile(file, ip, userid, filename, id, retention) result = "https://xygt.cc/{}".format(result) return result, status elif 'url' in request.form: url = request.form['url'] result, status = worker.shortURL(url, ip, userid, id, retention) @app.route('/') def getData(id): # Does it exist in the files DB? if Config.files.find_one({"id": id}) is not None: data = Config.files.find_one({"id": id}) with open(os.path.join(Config.fileDir, secure_filename(id)), "rb") as f: file = f.read() # Get MIME type from file, if fails then use magic try: mimetype = data["mimetype"] except KeyError: mimetype = magic.from_buffer(file, mime=True) # Return the file with the correct MIME type return send_file(io.BytesIO(file), mimetype=mimetype) # If not then check the URL Shortening DB elif Config.url.find_one({"id": id}) is not None: data = Config.url.find_one({"id": id}) return redirect(data["url"]) else: return random.choice(Errors.file404) @app.route('//info') def getInfo(id): return worker.idInfo(id) @app.route('/teapot') def teapot(): return 'I\'m a teapot. 418.', 418 @app.route('/register', methods=["GET", "POST"]) def register(): if current_user.is_authenticated: return redirect("/") else: if request.method == "GET": return render_template("register.html", form=RegistrationForm()) elif request.method == "POST": username = request.form.get("username") password = request.form.get("password") res = worker.registerUser(username, password) if res == True: flash("Successfully registered!", "success") return redirect("/login") else: flash("Something went wrong, sorry.", "danger") return redirect("/register") @app.route('/login', methods=["GET", "POST"]) def login(): if current_user.is_authenticated: return redirect("/") else: if request.method == "GET": return render_template("login.html", form=LoginForm()) elif request.method == "POST": username = request.form.get("username") password = request.form.get("password") userid = Config.users.find_one({"user": username})["userid"] user = User.get(userid) if user and bcrypt.check_password_hash(user.password, password): login_user(user) print(current_user.is_authenticated) flash("Successfully logged in!", "success") return redirect("/") else: flash("Incorrect username or password.", "danger") return redirect("/login") @app.route('/logout') def logout(): logout_user() return redirect("/")