FastAPI JWT Authentication: A Secure Example

by Jhon Lennon 45 views

Hey guys! Today, we're diving deep into the world of FastAPI and JWT (JSON Web Tokens) to build a secure authentication system. If you've ever wondered how to protect your API endpoints and ensure only authorized users can access them, you're in the right place. We'll walk through a practical example, breaking down each step to make it super easy to understand. So, buckle up and let's get started!

Understanding JWT Authentication

Before we jump into the code, let's quickly understand what JWT authentication is and why it's so popular. JWT is a standard for creating data with optional signature and/or encryption whose payload holds JSON that asserts a number of claims. In simpler terms, it's a way to securely transmit information between parties as a JSON object. This information can be verified and trusted because it is digitally signed. JWTs are commonly used for authentication because they allow you to verify the identity of users without needing to query a database repeatedly.

The process generally works like this: A user logs in with their credentials. The server verifies these credentials and, if valid, generates a JWT. This JWT is then sent back to the client. The client stores this JWT (usually in local storage or a cookie) and includes it in the headers of subsequent requests to the server. The server then verifies the JWT on each request to ensure the user is authenticated. The beauty of JWTs is that they are self-contained, meaning they contain all the necessary information to verify the user's identity.

Why is this so useful? Well, imagine you have a complex application with multiple microservices. Instead of each service needing to authenticate the user separately, they can all simply verify the JWT. This simplifies the authentication process and reduces the load on your authentication server. Plus, JWTs can include additional information about the user, such as their roles and permissions, allowing you to implement fine-grained access control.

Setting Up Your FastAPI Project

Alright, let's get our hands dirty and set up a FastAPI project. First things first, make sure you have Python installed. I recommend using Python 3.7 or higher. Once you have Python, you'll need to install FastAPI and its dependencies. Open your terminal and run the following command:

pip install fastapi uvicorn python-jose passlib bcrypt

Let's break down what each of these packages does:

  • fastapi: This is the web framework we'll be using.
  • uvicorn: This is an ASGI server that we'll use to run our FastAPI application.
  • python-jose: This library is used for encoding, decoding, and signing JWTs.
  • passlib: This provides secure password hashing.
  • bcrypt: This is a password hashing algorithm that we'll use with Passlib.

Once you have these packages installed, create a new directory for your project and create a file named main.py inside it. This is where we'll write our FastAPI code. Now, let's start building our authentication system.

Building the Authentication System

We'll start by defining our user model and the authentication endpoints. Here's a basic user model:

from pydantic import BaseModel

class User(BaseModel):
    username: str
    password: str

This is a simple Pydantic model that represents a user. It has two fields: username and password. Now, let's define our authentication endpoints. We'll need two endpoints: one for registering new users and one for logging in existing users. Here's how you can define these endpoints in FastAPI:

from fastapi import FastAPI, HTTPException
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
from passlib.context import CryptContext
from datetime import datetime, timedelta
from typing import Union, Any
from jose import JWTError, jwt


app = FastAPI()


# In-memory user database (replace with a real database in production)
users = {}

# Password hashing
crypt_context = CryptContext(schemes=["bcrypt"], deprecated="auto")

# JWT settings
SECRET_KEY = "your-secret-key"  # Change this in production!
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30


# Utility functions
def get_password_hash(password):
    return crypt_context.hash(password)


def verify_password(password, hashed_password):
    return crypt_context.verify(password, hashed_password)


def create_access_token(data: dict, expires_delta: Union[timedelta, None] = None):
    to_encode = data.copy()
    if expires_delta:
        expire = datetime.utcnow() + expires_delta
    else:
        expire = datetime.utcnow() + timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, SECRET_KEY, algorithm=ALGORITHM)
    return encoded_jwt


@app.post("/register")
async def register(user: User):
    if user.username in users:
        raise HTTPException(status_code=400, detail="Username already registered")
    hashed_password = get_password_hash(user.password)
    users[user.username] = {"password": hashed_password}
    return {"message": "User registered successfully"}


oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")

@app.post("/token")
async def login(form_data: OAuth2PasswordRequestForm):
    user = users.get(form_data.username)
    if not user:
        raise HTTPException(status_code=400, detail="Incorrect username or password")
    if not verify_password(form_data.password, user["password"]):
        raise HTTPException(status_code=400, detail="Incorrect username or password")

    access_token_expires = timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    access_token = create_access_token(
        data={"sub": form_data.username},
        expires_delta=access_token_expires
    )
    return {"access_token": access_token, "token_type": "bearer"}


async def get_current_user(token: str = Depends(oauth2_scheme)):
    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        username: str = payload.get("sub")
        if username is None:
            raise credentials_exception
        token_data = TokenData(username=username)
    except JWTError:
        raise credentials_exception
    user = get_user(username=token_data.username)
    if user is None:
        raise credentials_exception
    return user


@app.get("/users/me", response_model=User)
async def read_users_me(current_user: User = Depends(get_current_user)):
    return current_user

In this code, we define two endpoints: /register and /token. The /register endpoint allows new users to register by providing a username and password. The password is then hashed using bcrypt and stored in our in-memory user database (in a real-world application, you'd want to use a proper database like PostgreSQL or MySQL). The /token endpoint allows existing users to log in by providing their username and password. If the credentials are valid, a JWT is generated and returned to the client. The JWT is signed using a secret key, which you should definitely change in production.

We also defined a get_current_user function, which is used to authenticate users based on the JWT they provide in the Authorization header. This function decodes the JWT, verifies its signature, and extracts the username from the payload. It then retrieves the user from the database and returns it. If the JWT is invalid or the user doesn't exist, an HTTPException is raised.

Protecting Your API Endpoints

Now that we have our authentication system set up, let's see how we can protect our API endpoints using JWT authentication. We can use the Depends function in FastAPI to inject the get_current_user function into our endpoints. This will ensure that only authenticated users can access those endpoints. Here's an example:

from fastapi import Depends
from fastapi import status
from pydantic import BaseModel

class TokenData(BaseModel):
    username: Union[str, None] = None


# Example endpoint that requires authentication
@app.get("/protected")
async def protected(current_user: User = Depends(get_current_user)):
    return {"message": f"Hello, {current_user.username}!"}



def get_user(username: str):
    if username in users:
        return User(username=username, password=users[username]["password"])
    return None

In this code, we define a /protected endpoint that requires authentication. The current_user parameter is injected using the Depends function. This means that the get_current_user function will be called before the endpoint is executed. If the user is authenticated, their information will be available in the current_user parameter. If the user is not authenticated, an HTTPException will be raised.

Testing Your Authentication System

Now that we've built our authentication system and protected our API endpoints, let's test it out. First, run your FastAPI application using Uvicorn:

uvicorn main:app --reload

This will start your application on http://127.0.0.1:8000. Now, you can use a tool like Postman or curl to test your endpoints. First, register a new user by sending a POST request to the /register endpoint with a JSON body containing the username and password. Then, log in by sending a POST request to the /token endpoint with the username and password. This will return a JWT.

Finally, access the /protected endpoint by sending a GET request with the Authorization header set to Bearer <your_jwt>. If everything is set up correctly, you should see a message saying "Hello, <your_username>!". If you try to access the /protected endpoint without the Authorization header, you should get an HTTPException with a status code of 401.

Enhancements and Best Practices

Our example is a great starting point, but there's always room for improvement. Here are some enhancements and best practices to consider:

  • Use a Real Database: Our example uses an in-memory user database, which is not suitable for production. Use a real database like PostgreSQL or MySQL.
  • Implement Refresh Tokens: JWTs have a limited lifespan. Implement refresh tokens to allow users to stay logged in for longer periods without needing to re-enter their credentials.
  • Add Role-Based Access Control: Implement role-based access control to restrict access to certain endpoints based on the user's role.
  • Use Environment Variables: Store sensitive information like your secret key in environment variables instead of hardcoding them in your code.
  • Implement Token Revocation: Allow users to revoke their JWTs in case they are compromised.
  • Regularly Rotate Your Secret Key: To enhance security, regularly rotate your secret key.

Conclusion

So, there you have it! We've built a secure authentication system using FastAPI and JWTs. We've covered everything from setting up your project to protecting your API endpoints. While our example is basic, it provides a solid foundation for building more complex authentication systems. Remember to always follow security best practices and keep your code up-to-date to protect your users' data.

Keep experimenting, keep learning, and I'll catch you in the next one. Peace out!