Initial Commit

Author: ashutoshkrris

.gitignore

@@ -0,0 +1,28 @@
+.DS_Store
+.env
+.flaskenv
+*.pyc
+*.pyo
+env/
+venv/
+.venv/
+env*
+dist/
+build/
+*.egg
+*.egg-info/
+.tox/
+.cache/
+.pytest_cache/
+.idea/
+docs/_build/
+.vscode
+instance/
+
+# Coverage reports
+htmlcov/
+.coverage
+.coverage.*
+*,cover
+
+migrations/

config.py

@@ -0,0 +1,39 @@
+from decouple import config
+
+DATABASE_URI = config("DATABASE_URL")
+if DATABASE_URI.startswith("postgres://"):
+    DATABASE_URI = DATABASE_URI.replace("postgres://", "postgresql://", 1)
+
+
+class Config(object):
+    DEBUG = False
+    TESTING = False
+    CSRF_ENABLED = True
+    SECRET_KEY = config("SECRET_KEY", default="guess-me")
+    SQLALCHEMY_DATABASE_URI = DATABASE_URI
+    SQLALCHEMY_TRACK_MODIFICATIONS = False
+    BCRYPT_LOG_ROUNDS = 13
+    WTF_CSRF_ENABLED = True
+    DEBUG_TB_ENABLED = False
+    DEBUG_TB_INTERCEPT_REDIRECTS = False
+    APP_NAME = config("APP_NAME")
+
+
+class DevelopmentConfig(Config):
+    DEVELOPMENT = True
+    DEBUG = True
+    WTF_CSRF_ENABLED = False
+    DEBUG_TB_ENABLED = True
+
+
+class TestingConfig(Config):
+    TESTING = True
+    DEBUG = True
+    SQLALCHEMY_DATABASE_URI = "sqlite:///testdb.sqlite"
+    BCRYPT_LOG_ROUNDS = 1
+    WTF_CSRF_ENABLED = False
+
+
+class ProductionConfig(Config):
+    DEBUG = False
+    DEBUG_TB_ENABLED = False

manage.py

@@ -0,0 +1,9 @@
+from flask.cli import FlaskGroup
+
+from src import app
+
+cli = FlaskGroup(app)
+
+
+if __name__ == "__main__":
+    cli()

requirements.txt

@@ -0,0 +1,24 @@
+alembic==1.12.1
+bcrypt==4.0.1
+blinker==1.7.0
+click==8.1.7
+colorama==0.4.6
+Flask==3.0.0
+Flask-Bcrypt==1.0.1
+Flask-Login==0.6.3
+Flask-Migrate==4.0.5
+Flask-SQLAlchemy==3.1.1
+Flask-WTF==1.2.1
+greenlet==3.0.1
+itsdangerous==2.1.2
+Jinja2==3.1.2
+Mako==1.3.0
+MarkupSafe==2.1.3
+pyotp==2.9.0
+pypng==0.20220715.0
+python-decouple==3.8
+qrcode==7.4.2
+SQLAlchemy==2.0.23
+typing_extensions==4.8.0
+Werkzeug==3.0.1
+WTForms==3.1.1

src/__init__.py

@@ -0,0 +1,32 @@
+from decouple import config
+from flask import Flask
+from flask_bcrypt import Bcrypt
+from flask_login import LoginManager # Add this line
+from flask_migrate import Migrate
+from flask_sqlalchemy import SQLAlchemy
+
+app = Flask(__name__)
+app.config.from_object(config("APP_SETTINGS"))
+
+bcrypt = Bcrypt(app)
+db = SQLAlchemy(app)
+migrate = Migrate(app, db)
+
+login_manager = LoginManager()  # Add this line
+login_manager.init_app(app)  # Add this line
+
+# Registering blueprints
+from src.accounts.views import accounts_bp
+from src.core.views import core_bp
+
+app.register_blueprint(accounts_bp)
+app.register_blueprint(core_bp)
+
+from src.accounts.models import User
+
+login_manager.login_view = "accounts.login"
+login_manager.login_message_category = "danger"
+
+@login_manager.user_loader
+def load_user(user_id):
+    return User.query.filter(User.id == int(user_id)).first()

src/accounts/__init__.py

@@ -0,0 +1,44 @@
+from flask_wtf import FlaskForm
+from wtforms import StringField, PasswordField
+from wtforms.validators import DataRequired, EqualTo, Length, InputRequired
+
+from src.accounts.models import User
+
+
+class RegisterForm(FlaskForm):
+    username = StringField(
+        "Username", validators=[DataRequired(), Length(min=6, max=40)]
+    )
+    password = PasswordField(
+        "Password", validators=[DataRequired(), Length(min=6, max=25)]
+    )
+    confirm = PasswordField(
+        "Repeat password",
+        validators=[
+            DataRequired(),
+            EqualTo("password", message="Passwords must match."),
+        ],
+    )
+
+    def validate(self, extra_validators):
+        initial_validation = super(RegisterForm, self).validate(extra_validators)
+        if not initial_validation:
+            return False
+        user = User.query.filter_by(username=self.username.data).first()
+        if user:
+            self.username.errors.append("Username already registered")
+            return False
+        if self.password.data != self.confirm.data:
+            self.password.errors.append("Passwords must match")
+            return False
+        return True
+
+
+class LoginForm(FlaskForm):
+    username = StringField("Username", validators=[DataRequired()])
+    password = PasswordField("Password", validators=[DataRequired()])
+
+
+class TwoFactorForm(FlaskForm):
+    otp = StringField('Enter OTP', validators=[
+                      InputRequired(), Length(min=6, max=6)])

src/accounts/forms.py

@@ -0,0 +1,37 @@
+from datetime import datetime
+
+import pyotp
+from flask_login import UserMixin
+
+from src import bcrypt, db
+from config import Config
+
+
+class User(UserMixin, db.Model):
+
+    __tablename__ = "users"
+
+    id = db.Column(db.Integer, primary_key=True)
+    username = db.Column(db.String, unique=True, nullable=False)
+    password = db.Column(db.String, nullable=False)
+    created_at = db.Column(db.DateTime, nullable=False)
+    is_two_factor_authentication_enabled = db.Column(
+        db.Boolean, nullable=False, default=False)
+    secret_token = db.Column(db.String, unique=True)
+
+    def __init__(self, username, password):
+        self.username = username
+        self.password = bcrypt.generate_password_hash(password)
+        self.created_at = datetime.now()
+        self.secret_token = pyotp.random_base32()
+
+    def get_authentication_setup_uri(self):
+        return pyotp.totp.TOTP(self.secret_token).provisioning_uri(
+            name=self.username, issuer_name=Config.APP_NAME)
+
+    def is_otp_valid(self, user_otp):
+        totp = pyotp.parse_uri(self.get_authentication_setup_uri())
+        return totp.verify(user_otp)
+
+    def __repr__(self):
+        return f"<user {self.username}>"

src/accounts/models.py

@@ -0,0 +1,114 @@
+from src.utils import get_b64encoded_qr_image
+from .forms import LoginForm, RegisterForm, TwoFactorForm
+from src.accounts.models import User
+from src import db, bcrypt
+from flask_login import current_user, login_required, login_user, logout_user
+from flask import Blueprint, flash, redirect, render_template, request, url_for
+
+accounts_bp = Blueprint("accounts", __name__)
+
+HOME_URL = "core.home"
+SETUP_2FA_URL = "accounts.setup_two_factor_auth"
+VERIFY_2FA_URL = "accounts.verify_two_factor_auth"
+
+
+@accounts_bp.route("/register", methods=["GET", "POST"])
+def register():
+    if current_user.is_authenticated:
+        if current_user.is_two_factor_authentication_enabled:
+            flash("You are already registered.", "info")
+            return redirect(url_for(HOME_URL))
+        else:
+            flash(
+                "You have not enabled 2-Factor Authentication. Please enable first to login.", "info")
+            return redirect(url_for(SETUP_2FA_URL))
+    form = RegisterForm(request.form)
+    if form.validate_on_submit():
+        try:
+            user = User(username=form.username.data, password=form.password.data)
+            db.session.add(user)
+            db.session.commit()
+
+            login_user(user)
+            flash("You are registered. You have to enable 2-Factor Authentication first to login.", "success")
+
+            return redirect(url_for(SETUP_2FA_URL))
+        except Exception:
+            db.session.rollback()
+            flash("Registration failed. Please try again.", "danger")
+
+    return render_template("accounts/register.html", form=form)
+
+
+@accounts_bp.route("/login", methods=["GET", "POST"])
+def login():
+    if current_user.is_authenticated:
+        if current_user.is_two_factor_authentication_enabled:
+            flash("You are already logged in.", "info")
+            return redirect(url_for(HOME_URL))
+        else:
+            flash(
+                "You have not enabled 2-Factor Authentication. Please enable first to login.", "info")
+            return redirect(url_for(SETUP_2FA_URL))
+        
+    form = LoginForm(request.form)
+    if form.validate_on_submit():
+        user = User.query.filter_by(username=form.username.data).first()
+        if user and bcrypt.check_password_hash(user.password, request.form["password"]):
+            login_user(user)
+            if not current_user.is_two_factor_authentication_enabled:
+                flash(
+                    "You have not enabled 2-Factor Authentication. Please enable first to login.", "info")
+                return redirect(url_for(SETUP_2FA_URL))
+            return redirect(url_for(VERIFY_2FA_URL))
+        elif not user:
+            flash("You are not registered. Please register.", "danger")
+        else:
+            flash("Invalid username and/or password.", "danger")
+    return render_template("accounts/login.html", form=form)
+
+
+@accounts_bp.route("/logout")
+@login_required
+def logout():
+    logout_user()
+    flash("You were logged out.", "success")
+    return redirect(url_for("accounts.login"))
+
+
+@accounts_bp.route("/setup-2fa")
+@login_required
+def setup_two_factor_auth():
+    secret = current_user.secret_token
+    uri = current_user.get_authentication_setup_uri()
+    base64_qr_image = get_b64encoded_qr_image(uri)
+    return render_template("accounts/setup-2fa.html", secret=secret, qr_image=base64_qr_image)
+
+
+@accounts_bp.route("/verify-2fa", methods=["GET", "POST"])
+@login_required
+def verify_two_factor_auth():
+    form = TwoFactorForm(request.form)
+    if form.validate_on_submit():
+        if current_user.is_otp_valid(form.otp.data):
+            if current_user.is_two_factor_authentication_enabled:
+                flash("2FA verification successful. You are logged in!", "success")
+                return redirect(url_for(HOME_URL))
+            else:
+                try:
+                    current_user.is_two_factor_authentication_enabled = True
+                    db.session.commit()
+                    flash("2FA setup successful. You are logged in!", "success")
+                    return redirect(url_for(HOME_URL))
+                except Exception:
+                    db.session.rollback()
+                    flash("2FA setup failed. Please try again.", "danger")
+                    return redirect(url_for(VERIFY_2FA_URL))
+        else:
+            flash("Invalid OTP. Please try again.", "danger")
+            return redirect(url_for(VERIFY_2FA_URL))
+    else:
+        if not current_user.is_two_factor_authentication_enabled:
+            flash(
+                "You have not enabled 2-Factor Authentication. Please enable it first.", "info")
+        return render_template("accounts/verify-2fa.html", form=form)