diff --git a/app/models.py b/app/models.py index 56dd11f..2432836 100644 --- a/app/models.py +++ b/app/models.py @@ -277,6 +277,10 @@ class User(Model): is_superuser = fields.BooleanField(default=False) user_class = fields.CharEnumField(enum_type=UserClassEnum, null=True, default=UserClassEnum.PIPELINE_USER) data_source = fields.ForeignKeyField("app.DataSource", related_name="users", null=True) + # 2FA + secret_key = fields.CharField(max_length=255, null=True) + is_2fa_required = fields.BooleanField(default=False) + is_2fa_activated = fields.BooleanField(default=False) created_at = fields.DatetimeField(auto_now_add=True) updated_at = fields.DatetimeField(auto_now=True) diff --git a/app/routers/auth.py b/app/routers/auth.py index 3e23581..f6d123f 100644 --- a/app/routers/auth.py +++ b/app/routers/auth.py @@ -1,43 +1,127 @@ # -*- coding: utf-8 -*- -from datetime import timedelta +import io from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, status from fastapi.security import OAuth2PasswordRequestForm +from fastapi.responses import StreamingResponse -from app import config from app.models import User from app.types.pydantic_models import Token -from app.utils import authenticate_user, create_access_token +from app.utils import authenticate_user, generate_user_token +from app.security import TwoFactorAuth +from app.dependencies import ( + get_current_frontend_user +) router = APIRouter(prefix="/auth", tags=["Autenticação"]) @router.post("/token") -async def login_for_access_token( - form_data: Annotated[OAuth2PasswordRequestForm, Depends()] +async def login_without_2fa( + form_data: Annotated[OAuth2PasswordRequestForm, Depends()], ) -> Token: - user: User = await authenticate_user(form_data.username, form_data.password) + user = await authenticate_user(form_data.username, form_data.password) + if not user: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + headers={"WWW-Authenticate": "Bearer"}, + ) + + if user.is_2fa_required: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="2FA required. Use the /2fa/login/ endpoint", + headers={"WWW-Authenticate": "Bearer"}, + ) + + return { + "access_token": generate_user_token(user), + "token_type": "bearer" + } + + +@router.post("/2fa/is-2fa-active/") +async def is_2fa_active( + form_data: Annotated[OAuth2PasswordRequestForm, Depends()], +) -> bool: + user = await authenticate_user(form_data.username, form_data.password) + if not user: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + headers={"WWW-Authenticate": "Bearer"}, + ) + + return user.is_2fa_activated + +@router.post("/2fa/login/") +async def login_with_2fa( + form_data: Annotated[OAuth2PasswordRequestForm, Depends()], + totp_code: str, +) -> Token: + + user = await authenticate_user(form_data.username, form_data.password) if not user: raise HTTPException( - status_code = status.HTTP_401_UNAUTHORIZED, - detail = "Incorrect username or password", - headers = {"WWW-Authenticate": "Bearer"}, + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + headers={"WWW-Authenticate": "Bearer"}, ) - access_token_expires = timedelta( - minutes = config.JWT_ACCESS_TOKEN_EXPIRE_MINUTES - ) + secret_key = await TwoFactorAuth.get_or_create_secret_key(user.id) + two_factor_auth = TwoFactorAuth(user.id, secret_key) - access_token = create_access_token( - data = {"sub": user.username}, - expires_delta = access_token_expires - ) + is_valid = two_factor_auth.verify_totp_code(totp_code) + if not is_valid: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect OTP", + headers={"WWW-Authenticate": "Bearer"}, + ) + if not user.is_2fa_activated: + user.is_2fa_activated = True + await user.save() return { - "access_token": access_token, - "token_type": "bearer" - } \ No newline at end of file + "access_token": generate_user_token(user), + "token_type": "bearer", + } + + +@router.post("/2fa/enable/") +async def enable_2fa( + current_user: Annotated[User, Depends(get_current_frontend_user)], +): + secret_key = await TwoFactorAuth.get_or_create_secret_key(current_user.id) + two_factor_auth = TwoFactorAuth(current_user.id, secret_key) + + return { + "secret_key": two_factor_auth.secret_key + } + + +@router.get("/2fa/generate-qrcode/") +async def generate_qrcode( + form_data: Annotated[OAuth2PasswordRequestForm, Depends()], +): + current_user = await authenticate_user(form_data.username, form_data.password) + if not current_user: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail="Incorrect username or password", + headers={"WWW-Authenticate": "Bearer"}, + ) + + secret_key = await TwoFactorAuth.get_or_create_secret_key(current_user.id) + two_factor_auth = TwoFactorAuth(current_user.id, secret_key) + + qr_code = two_factor_auth.qr_code + if qr_code is None: + raise HTTPException(status_code=404, detail="User not found") + + return StreamingResponse(io.BytesIO(qr_code), media_type="image/png") diff --git a/app/security.py b/app/security.py new file mode 100644 index 0000000..b3e2298 --- /dev/null +++ b/app/security.py @@ -0,0 +1,77 @@ +# -*- coding: utf-8 -*- +import base64 +import io +import secrets +from typing import Optional + +import qrcode +from pyotp import TOTP + +from app.models import User + + +class TwoFactorAuth: + + def __init__(self, user_id: str, secret_key: str): + self._user_id = user_id + self._secret_key = secret_key + self._totp = TOTP(self._secret_key) + self._qr_cache: Optional[bytes] = None + + @property + def totp(self) -> TOTP: + return self._totp + + @property + def secret_key(self) -> str: + return self._secret_key + + @staticmethod + def _generate_secret_key() -> str: + secret_bytes = secrets.token_bytes(20) + secret_key = base64.b32encode(secret_bytes).decode("utf-8") + return secret_key + + @staticmethod + async def get_or_create_secret_key(user_id: str) -> str: + user = await User.get_or_none(id=user_id) + + if not user: + raise ValueError(f"User with id {user_id} not found") + + # If User doesn't have a secret_key, create one + if not user.secret_key: + secret_key = TwoFactorAuth._generate_secret_key() + user.secret_key = secret_key + await user.save() + + return user.secret_key + + def _create_qr_code(self) -> bytes: + uri = self.totp.provisioning_uri( + name=str(self._user_id), + issuer_name="2FA", + ) + img = qrcode.make(uri) + img_byte_array = io.BytesIO() + img.save(img_byte_array) + img_byte_array.seek(0) + return img_byte_array.getvalue() + + @property + def qr_code(self) -> bytes: + if self._qr_cache is None: + self._qr_cache = self._create_qr_code() + return self._qr_cache + + def verify_totp_code(self, totp_code: str) -> bool: + return self.totp.verify(totp_code) + + +async def get_two_factor_auth( + user_id: str +) -> TwoFactorAuth: + secret_key = await TwoFactorAuth.get_or_create_secret_key( + user_id + ) + return TwoFactorAuth(user_id, secret_key) \ No newline at end of file diff --git a/app/utils.py b/app/utils.py index 67be48f..25aa6cf 100644 --- a/app/utils.py +++ b/app/utils.py @@ -18,6 +18,15 @@ pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto") +def generate_user_token(user: User) -> str: + access_token_expires = timedelta(minutes=config.JWT_ACCESS_TOKEN_EXPIRE_MINUTES) + + access_token = create_access_token( + data={"sub": user.username}, expires_delta=access_token_expires + ) + + return access_token + async def authenticate_user(username: str, password: str) -> User: """Authenticate a user. diff --git a/migrations/app/27_20240830145823_update.py b/migrations/app/27_20240830145823_update.py new file mode 100644 index 0000000..f17dac5 --- /dev/null +++ b/migrations/app/27_20240830145823_update.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +from tortoise import BaseDBAsyncClient + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" ADD "secret_key" VARCHAR(255);""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" DROP COLUMN "secret_key";""" diff --git a/migrations/app/28_20240830163720_update.py b/migrations/app/28_20240830163720_update.py new file mode 100644 index 0000000..0070d08 --- /dev/null +++ b/migrations/app/28_20240830163720_update.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +from tortoise import BaseDBAsyncClient + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" ADD "is_2fa_enabled" BOOL NOT NULL DEFAULT False;""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" DROP COLUMN "is_2fa_enabled";""" diff --git a/migrations/app/29_20240830164032_update.py b/migrations/app/29_20240830164032_update.py new file mode 100644 index 0000000..a54b4d4 --- /dev/null +++ b/migrations/app/29_20240830164032_update.py @@ -0,0 +1,12 @@ +# -*- coding: utf-8 -*- +from tortoise import BaseDBAsyncClient + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" ADD "is_2fa_active" BOOL NOT NULL DEFAULT False;""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" DROP COLUMN "is_2fa_active";""" diff --git a/migrations/app/30_20240830164307_update.py b/migrations/app/30_20240830164307_update.py new file mode 100644 index 0000000..0c4c781 --- /dev/null +++ b/migrations/app/30_20240830164307_update.py @@ -0,0 +1,16 @@ +# -*- coding: utf-8 -*- +from tortoise import BaseDBAsyncClient + + +async def upgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" RENAME COLUMN "is_2fa_enabled" TO "is_2fa_required"; + ALTER TABLE "user" RENAME COLUMN "is_2fa_active" TO "is_2fa_activated";""" + + +async def downgrade(db: BaseDBAsyncClient) -> str: + return """ + ALTER TABLE "user" RENAME COLUMN "is_2fa_required" TO "is_2fa_active"; + ALTER TABLE "user" RENAME COLUMN "is_2fa_activated" TO "is_2fa_active"; + ALTER TABLE "user" RENAME COLUMN "is_2fa_required" TO "is_2fa_enabled"; + ALTER TABLE "user" RENAME COLUMN "is_2fa_activated" TO "is_2fa_enabled";""" diff --git a/poetry.lock b/poetry.lock index bed4472..7dad60e 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1938,6 +1938,20 @@ cffi = ">=1.4.1" docs = ["sphinx (>=1.6.5)", "sphinx-rtd-theme"] tests = ["hypothesis (>=3.27.0)", "pytest (>=3.2.1,!=3.3.0)"] +[[package]] +name = "pyotp" +version = "2.9.0" +description = "Python One Time Password Library" +optional = false +python-versions = ">=3.7" +files = [ + {file = "pyotp-2.9.0-py3-none-any.whl", hash = "sha256:81c2e5865b8ac55e825b0358e496e1d9387c811e85bb40e71a3b29b288963612"}, + {file = "pyotp-2.9.0.tar.gz", hash = "sha256:346b6642e0dbdde3b4ff5a930b664ca82abfa116356ed48cc42c7d6590d36f63"}, +] + +[package.extras] +test = ["coverage", "mypy", "ruff", "wheel"] + [[package]] name = "pyparsing" version = "3.1.2" @@ -1963,6 +1977,17 @@ files = [ {file = "pypika_tortoise-0.1.6-py3-none-any.whl", hash = "sha256:2d68bbb7e377673743cff42aa1059f3a80228d411fbcae591e4465e173109fd8"}, ] +[[package]] +name = "pypng" +version = "0.20220715.0" +description = "Pure Python library for saving and loading PNG images" +optional = false +python-versions = "*" +files = [ + {file = "pypng-0.20220715.0-py3-none-any.whl", hash = "sha256:4a43e969b8f5aaafb2a415536c1a8ec7e341cd6a3f957fd5b5f32a4cfeed902c"}, + {file = "pypng-0.20220715.0.tar.gz", hash = "sha256:739c433ba96f078315de54c0db975aee537cbc3e1d0ae4ed9aab0ca1e427e2c1"}, +] + [[package]] name = "pytest" version = "7.4.4" @@ -2146,6 +2171,29 @@ files = [ {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, ] +[[package]] +name = "qrcode" +version = "7.4.2" +description = "QR Code image generator" +optional = false +python-versions = ">=3.7" +files = [ + {file = "qrcode-7.4.2-py3-none-any.whl", hash = "sha256:581dca7a029bcb2deef5d01068e39093e80ef00b4a61098a2182eac59d01643a"}, + {file = "qrcode-7.4.2.tar.gz", hash = "sha256:9dd969454827e127dbd93696b20747239e6d540e082937c90f14ac95b30f5845"}, +] + +[package.dependencies] +colorama = {version = "*", markers = "platform_system == \"Windows\""} +pypng = "*" +typing-extensions = "*" + +[package.extras] +all = ["pillow (>=9.1.0)", "pytest", "pytest-cov", "tox", "zest.releaser[recommended]"] +dev = ["pytest", "pytest-cov", "tox"] +maintainer = ["zest.releaser[recommended]"] +pil = ["pillow (>=9.1.0)"] +test = ["coverage", "pytest"] + [[package]] name = "regex" version = "2024.7.24" @@ -2807,4 +2855,4 @@ dev = ["black (>=19.3b0)", "pytest (>=4.6.2)"] [metadata] lock-version = "2.0" python-versions = "^3.11" -content-hash = "1ec4944b80ec680487b4e0ffc0144b035cc4fc7efeb12e81f5b585127b8c0271" +content-hash = "5715de4721d690e2ecff456af1ed2d255f810d0c660d5cb7f26bd307d92e4199" diff --git a/pyproject.toml b/pyproject.toml index 37f0d01..f9022a8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,8 @@ idna = "3.7" basedosdados = "^2.0.0b16" nltk = "^3.9.1" asyncer = "^0.0.8" +qrcode = "^7.4.2" +pyotp = "^2.9.0" [tool.poetry.group.dev.dependencies]