stable-diffusion.cpp-rest/include/jwt_auth.h

#ifndef JWT_AUTH_H
#define JWT_AUTH_H

#include <string>
#include <vector>
#include <map>
#include <chrono>
#include <memory>

/**
 * @brief JWT (JSON Web Token) authentication implementation
 *
 * This class provides JWT token generation and validation functionality
 * for the stable-diffusion.cpp-rest server. It supports HS256 algorithm
 * and includes claims for user identification, expiration, and roles.
 */
class JWTAuth {
public:
    /**
     * @brief JWT token claims structure
     */
    struct Claims {
        std::string userId;        ///< User identifier
        std::string username;      ///< Username
        std::string role;          ///< User role (admin, user, etc.)
        std::vector<std::string> permissions;  ///< User permissions
        int64_t issuedAt;         ///< Issued at timestamp
        int64_t expiresAt;        ///< Expiration timestamp
        std::string issuer;        ///< Token issuer
        std::string audience;      ///< Token audience
    };

    /**
     * @brief Authentication result structure
     */
    struct AuthResult {
        bool success;              ///< Authentication success status
        std::string userId;        ///< User ID if successful
        std::string username;      ///< Username if successful
        std::string role;          ///< User role if successful
        std::vector<std::string> permissions;  ///< Permissions if successful
        std::string errorMessage;  ///< Error message if failed
        std::string errorCode;     ///< Error code for API responses
    };

    /**
     * @brief Construct a new JWT Auth object
     *
     * @param secret Secret key for signing tokens
     * @param tokenExpirationMinutes Token expiration time in minutes (default: 60)
     * @param issuer Token issuer (default: "stable-diffusion-rest")
     * @param audience Token audience (default: "stable-diffusion-rest")
     */
    explicit JWTAuth(const std::string& secret,
                    int tokenExpirationMinutes = 60,
                    const std::string& issuer = "stable-diffusion-rest",
                    const std::string& audience = "stable-diffusion-rest");

    /**
     * @brief Destroy the JWT Auth object
     */
    ~JWTAuth();

    /**
     * @brief Generate a JWT token for the given user
     *
     * @param userId User identifier
     * @param username Username
     * @param role User role
     * @param permissions User permissions list
     * @return std::string JWT token string, empty on failure
     */
    std::string generateToken(const std::string& userId,
                             const std::string& username,
                             const std::string& role,
                             const std::vector<std::string>& permissions = {});

    /**
     * @brief Validate a JWT token and extract claims
     *
     * @param token JWT token string
     * @return AuthResult Authentication result with user information
     */
    AuthResult validateToken(const std::string& token);

    /**
     * @brief Refresh an existing token (extend expiration)
     *
     * @param token Existing JWT token
     * @return std::string New JWT token, empty on failure
     */
    std::string refreshToken(const std::string& token);

    /**
     * @brief Extract token from Authorization header
     *
     * @param authHeader Authorization header value
     * @return std::string Token string, empty if not found or invalid format
     */
    static std::string extractTokenFromHeader(const std::string& authHeader);

    /**
     * @brief Check if user has required permission
     *
     * @param permissions User permissions list
     * @param requiredPermission Required permission to check
     * @return true if user has permission, false otherwise
     */
    static bool hasPermission(const std::vector<std::string>& permissions,
                             const std::string& requiredPermission);

    /**
     * @brief Check if user has any of the required permissions
     *
     * @param permissions User permissions list
     * @param requiredPermissions List of permissions to check (any one is sufficient)
     * @return true if user has any of the permissions, false otherwise
     */
    static bool hasAnyPermission(const std::vector<std::string>& permissions,
                                const std::vector<std::string>& requiredPermissions);

    /**
     * @brief Get token expiration time
     *
     * @param token JWT token string
     * @return int64_t Expiration timestamp, 0 on failure
     */
    int64_t getTokenExpiration(const std::string& token);

    /**
     * @brief Check if token is expired
     *
     * @param token JWT token string
     * @return true if token is expired, false otherwise
     */
    bool isTokenExpired(const std::string& token);

    /**
     * @brief Set token expiration time
     *
     * @param minutes Expiration time in minutes
     */
    void setTokenExpiration(int minutes);

    /**
     * @brief Get token expiration time in minutes
     *
     * @return int Token expiration time in minutes
     */
    int getTokenExpiration() const;

    /**
     * @brief Set issuer for tokens
     *
     * @param issuer Issuer string
     */
    void setIssuer(const std::string& issuer);

    /**
     * @brief Get issuer string
     *
     * @return std::string Issuer string
     */
    std::string getIssuer() const;

    /**
     * @brief Generate a random API key
     *
     * @param length Length of the API key (default: 32)
     * @return std::string Random API key
     */
    static std::string generateApiKey(int length = 32);

    /**
     * @brief Validate API key format
     *
     * @param apiKey API key string
     * @return true if format is valid, false otherwise
     */
    static bool validateApiKeyFormat(const std::string& apiKey);

private:
    std::string m_secret;                    ///< Secret key for signing
    int m_tokenExpirationMinutes;            ///< Token expiration in minutes
    std::string m_issuer;                    ///< Token issuer
    std::string m_audience;                  ///< Token audience

    /**
     * @brief Base64 URL encode a string
     *
     * @param input Input string
     * @return std::string Base64 URL encoded string
     */
    static std::string base64UrlEncode(const std::string& input);

    /**
     * @brief Base64 URL decode a string
     *
     * @param input Base64 URL encoded string
     * @return std::string Decoded string
     */
    static std::string base64UrlDecode(const std::string& input);

    /**
     * @brief Create JWT header
     *
     * @return std::string JWT header JSON string
     */
    std::string createHeader() const;

    /**
     * @brief Create JWT payload from claims
     *
     * @param claims Token claims
     * @return std::string JWT payload JSON string
     */
    std::string createPayload(const Claims& claims) const;

    /**
     * @brief Parse JWT payload from token
     *
     * @param token JWT token
     * @return Claims Parsed claims, empty on failure
     */
    Claims parsePayload(const std::string& token) const;

    /**
     * @brief Create HMAC-SHA256 signature
     *
     * @param header Payload header
     * @param payload Payload data
     * @return std::string Signature string
     */
    std::string createSignature(const std::string& header, const std::string& payload) const;

    /**
     * @brief Verify HMAC-SHA256 signature
     *
     * @param header Payload header
     * @param payload Payload data
     * @param signature Signature to verify
     * @return true if signature is valid, false otherwise
     */
    bool verifySignature(const std::string& header, const std::string& payload, const std::string& signature) const;

    /**
     * @brief Split JWT token into parts
     *
     * @param token JWT token
     * @return std::vector<std::string> Token parts (header, payload, signature)
     */
    static std::vector<std::string> splitToken(const std::string& token);

    /**
     * @brief Get current timestamp in seconds
     *
     * @return int64_t Current timestamp
     */
    static int64_t getCurrentTimestamp();

    /**
     * @brief Generate random string
     *
     * @param length Length of the string
     * @return std::string Random string
     */
    static std::string generateRandomString(int length);
};

#endif // JWT_AUTH_H