wind_power_cal/backend/src/auth/AuthManager.cpp

#include "AuthManager.h"

#include <trantor/utils/Logger.h>

#include <ctime>
#include <filesystem>
#include <regex>
#include <string>

#include <drogon/utils/Utilities.h>
#include <sqlite3.h>

namespace {

namespace fs = std::filesystem;
using json = nlohmann::json;

constexpr int kPasswordHashRounds = 20000;
constexpr auto kSessionLifetime = std::chrono::hours(12);
const fs::path kDatabasePath = fs::path("data") / "wind_power.db";

struct AccountRecord {
    std::string username;
    std::string password_salt;
    std::string password_hash;
    bool is_admin = false;
    std::string expires_on;
    std::string created_at;
    bool enabled = true;
};

class Statement {
public:
    Statement(sqlite3* database, const char* sql) {
        if (sqlite3_prepare_v2(database, sql, -1, &statement_, nullptr) != SQLITE_OK) {
            statement_ = nullptr;
        }
    }

    ~Statement() {
        if (statement_ != nullptr) {
            sqlite3_finalize(statement_);
        }
    }

    Statement(const Statement&) = delete;
    Statement& operator=(const Statement&) = delete;

    sqlite3_stmt* Get() const {
        return statement_;
    }

    explicit operator bool() const {
        return statement_ != nullptr;
    }

private:
    sqlite3_stmt* statement_ = nullptr;
};

std::string ColumnText(sqlite3_stmt* statement, int column) {
    const auto* value = sqlite3_column_text(statement, column);
    return value == nullptr ? "" : reinterpret_cast<const char*>(value);
}

std::string CurrentDate() {
    const auto now = std::time(nullptr);
    std::tm local_time = {};
    localtime_r(&now, &local_time);
    char buffer[11] = {};
    std::strftime(buffer, sizeof(buffer), "%Y-%m-%d", &local_time);
    return buffer;
}

std::string CurrentDateTime() {
    const auto now = std::time(nullptr);
    std::tm local_time = {};
    localtime_r(&now, &local_time);
    char buffer[20] = {};
    std::strftime(buffer, sizeof(buffer), "%Y-%m-%d %H:%M:%S", &local_time);
    return buffer;
}

bool IsValidDate(const std::string& value) {
    if (!std::regex_match(value, std::regex(R"(^\d{4}-\d{2}-\d{2}$)"))) {
        return false;
    }
    std::tm parsed = {};
    if (strptime(value.c_str(), "%Y-%m-%d", &parsed) == nullptr) {
        return false;
    }
    parsed.tm_isdst = -1;
    const auto timestamp = std::mktime(&parsed);
    if (timestamp == static_cast<std::time_t>(-1)) {
        return false;
    }
    std::tm normalized = {};
    localtime_r(&timestamp, &normalized);
    char buffer[11] = {};
    std::strftime(buffer, sizeof(buffer), "%Y-%m-%d", &normalized);
    return value == buffer;
}

bool IsExpired(const AccountRecord& account) {
    return !account.expires_on.empty() && CurrentDate() > account.expires_on;
}

std::string HashPassword(const std::string& password, const std::string& salt) {
    std::string digest = password + ":" + salt;
    for (int round = 0; round < kPasswordHashRounds; ++round) {
        digest = drogon::utils::getSha256(salt + ":" + digest);
    }
    return digest;
}

bool ConstantTimeEquals(const std::string& left, const std::string& right) {
    if (left.size() != right.size()) {
        return false;
    }
    unsigned char difference = 0;
    for (size_t index = 0; index < left.size(); ++index) {
        difference |= static_cast<unsigned char>(left[index] ^ right[index]);
    }
    return difference == 0;
}

bool ValidateUsername(const std::string& username, std::string& error) {
    if (!std::regex_match(username, std::regex(R"(^[A-Za-z0-9_.-]{3,32}$)"))) {
        error = "账号只能包含字母、数字、点、下划线或短横线，长度为 3-32 位";
        return false;
    }
    return true;
}

bool ValidatePassword(const std::string& password, std::string& error) {
    if (password.size() < 6 || password.size() > 128) {
        error = "密码长度必须为 6-128 位";
        return false;
    }
    return true;
}

AuthUser ToAuthUser(const AccountRecord& account) {
    return AuthUser{
        account.username,
        account.is_admin,
        account.expires_on,
        account.enabled,
    };
}

std::optional<AccountRecord> FindAccount(sqlite3* database, const std::string& username) {
    Statement statement(database,
        "SELECT username, password_salt, password_hash, is_admin, expires_on, created_at, "
        "enabled "
        "FROM accounts WHERE username = ?");
    if (!statement) {
        return std::nullopt;
    }
    sqlite3_bind_text(statement.Get(), 1, username.c_str(), -1, SQLITE_TRANSIENT);
    if (sqlite3_step(statement.Get()) != SQLITE_ROW) {
        return std::nullopt;
    }
    return AccountRecord{
        ColumnText(statement.Get(), 0),
        ColumnText(statement.Get(), 1),
        ColumnText(statement.Get(), 2),
        sqlite3_column_int(statement.Get(), 3) != 0,
        ColumnText(statement.Get(), 4),
        ColumnText(statement.Get(), 5),
        sqlite3_column_int(statement.Get(), 6) != 0,
    };
}

bool HasColumn(sqlite3* database, const std::string& table, const std::string& column) {
    Statement statement(database, ("PRAGMA table_info(" + table + ")").c_str());
    if (!statement) {
        return false;
    }
    while (sqlite3_step(statement.Get()) == SQLITE_ROW) {
        if (ColumnText(statement.Get(), 1) == column) {
            return true;
        }
    }
    return false;
}

}  // namespace

AuthManager& AuthManager::Instance() {
    static AuthManager instance;
    return instance;
}

AuthManager::~AuthManager() {
    if (database_ != nullptr) {
        sqlite3_close(database_);
    }
}

bool AuthManager::Initialize() {
    std::lock_guard<std::mutex> lock(mutex_);
    return OpenDatabaseLocked() && CreateSchemaLocked() && BootstrapAdminLocked();
}

bool AuthManager::OpenDatabaseLocked() {
    if (database_ != nullptr) {
        return true;
    }
    fs::create_directories(kDatabasePath.parent_path());
    if (sqlite3_open_v2(kDatabasePath.c_str(),
                        &database_,
                        SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX,
                        nullptr) != SQLITE_OK) {
        return false;
    }
    sqlite3_busy_timeout(database_, 5000);
    return sqlite3_exec(database_, "PRAGMA journal_mode=WAL;", nullptr, nullptr, nullptr) ==
        SQLITE_OK;
}

bool AuthManager::CreateSchemaLocked() {
    constexpr const char* kSchema =
        "CREATE TABLE IF NOT EXISTS accounts ("
        "username TEXT PRIMARY KEY NOT NULL,"
        "password_salt TEXT NOT NULL,"
        "password_hash TEXT NOT NULL,"
        "is_admin INTEGER NOT NULL DEFAULT 0,"
        "expires_on TEXT NOT NULL DEFAULT '',"
        "created_at TEXT NOT NULL,"
        "enabled INTEGER NOT NULL DEFAULT 1"
        ");";
    if (sqlite3_exec(database_, kSchema, nullptr, nullptr, nullptr) != SQLITE_OK) {
        return false;
    }
    if (HasColumn(database_, "accounts", "enabled")) {
        return true;
    }
    constexpr const char* kAddEnabled =
        "ALTER TABLE accounts ADD COLUMN enabled INTEGER NOT NULL DEFAULT 1";
    return sqlite3_exec(database_, kAddEnabled, nullptr, nullptr, nullptr) == SQLITE_OK;
}

bool AuthManager::BootstrapAdminLocked() {
    if (FindAccount(database_, "admin").has_value()) {
        return true;
    }
    const auto salt = drogon::utils::secureRandomString(32);
    const auto password_hash = HashPassword("nwl888888", salt);
    const auto created_at = CurrentDateTime();
    Statement statement(database_,
        "INSERT INTO accounts "
        "(username, password_salt, password_hash, is_admin, expires_on, created_at) "
        "VALUES (?, ?, ?, 1, '', ?)");
    if (!statement) {
        return false;
    }
    sqlite3_bind_text(statement.Get(), 1, "admin", -1, SQLITE_STATIC);
    sqlite3_bind_text(statement.Get(), 2, salt.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 3, password_hash.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 4, created_at.c_str(), -1, SQLITE_TRANSIENT);
    return sqlite3_step(statement.Get()) == SQLITE_DONE;
}

AuthResult AuthManager::Login(const std::string& username,
                              const std::string& password,
                              std::string& token) {
    std::lock_guard<std::mutex> lock(mutex_);
    const auto account = FindAccount(database_, username);
    if (!account.has_value() ||
        !ConstantTimeEquals(HashPassword(password, account->password_salt),
                            account->password_hash)) {
        return {};
    }
    if (!account->enabled) {
        return {AuthStatus::kDisabled, ToAuthUser(account.value())};
    }
    if (IsExpired(account.value())) {
        return {AuthStatus::kExpired, ToAuthUser(account.value())};
    }
    token = drogon::utils::secureRandomString(48);
    sessions_[token] = Session{username, std::chrono::system_clock::now() + kSessionLifetime};
    return {AuthStatus::kAuthenticated, ToAuthUser(account.value())};
}

AuthResult AuthManager::Authenticate(const std::string& token) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (token.empty()) {
        return {};
    }
    const auto session = sessions_.find(token);
    if (session == sessions_.end() || session->second.expires_at <= std::chrono::system_clock::now()) {
        if (session != sessions_.end()) {
            sessions_.erase(session);
        }
        return {};
    }
    const auto account = FindAccount(database_, session->second.username);
    if (!account.has_value()) {
        sessions_.erase(session);
        return {};
    }
    if (!account->enabled) {
        sessions_.erase(session);
        return {AuthStatus::kDisabled, ToAuthUser(account.value())};
    }
    if (IsExpired(account.value())) {
        sessions_.erase(session);
        return {AuthStatus::kExpired, ToAuthUser(account.value())};
    }
    return {AuthStatus::kAuthenticated, ToAuthUser(account.value())};
}

void AuthManager::Logout(const std::string& token) {
    std::lock_guard<std::mutex> lock(mutex_);
    sessions_.erase(token);
}

nlohmann::json AuthManager::ListAccounts() {
    std::lock_guard<std::mutex> lock(mutex_);
    json result = json::array();
    Statement statement(database_,
        "SELECT username, is_admin, expires_on, created_at, enabled "
        "FROM accounts ORDER BY is_admin DESC, username");
    if (!statement) {
        return result;
    }
    while (sqlite3_step(statement.Get()) == SQLITE_ROW) {
        AccountRecord account;
        account.username = ColumnText(statement.Get(), 0);
        account.is_admin = sqlite3_column_int(statement.Get(), 1) != 0;
        account.expires_on = ColumnText(statement.Get(), 2);
        account.created_at = ColumnText(statement.Get(), 3);
        account.enabled = sqlite3_column_int(statement.Get(), 4) != 0;
        result.push_back({
            {"username", account.username},
            {"is_admin", account.is_admin},
            {"expires_on", account.expires_on},
            {"permanent", account.expires_on.empty()},
            {"enabled", account.enabled},
            {"created_at", account.created_at},
            {"expired", IsExpired(account)},
        });
    }
    return result;
}

bool AuthManager::CreateAccount(const std::string& username,
                                const std::string& password,
                                const std::string& expires_on,
                                std::string& error) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (!ValidateUsername(username, error) || !ValidatePassword(password, error)) {
        return false;
    }
    if (!expires_on.empty() &&
        (!IsValidDate(expires_on) || expires_on < CurrentDate())) {
        error = "有效期必须是今天或之后的有效日期";
        return false;
    }
    if (FindAccount(database_, username).has_value()) {
        error = "账号已存在";
        return false;
    }
    const auto salt = drogon::utils::secureRandomString(32);
    const auto password_hash = HashPassword(password, salt);
    const auto created_at = CurrentDateTime();
    Statement statement(database_,
        "INSERT INTO accounts "
        "(username, password_salt, password_hash, is_admin, expires_on, created_at) "
        "VALUES (?, ?, ?, 0, ?, ?)");
    if (!statement) {
        error = "保存账号失败";
        return false;
    }
    sqlite3_bind_text(statement.Get(), 1, username.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 2, salt.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 3, password_hash.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 4, expires_on.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 5, created_at.c_str(), -1, SQLITE_TRANSIENT);
    if (sqlite3_step(statement.Get()) != SQLITE_DONE) {
        error = "保存账号失败";
        return false;
    }
    return true;
}

bool AuthManager::UpdateAccount(const std::string& username,
                                const std::optional<std::string>& password,
                                const std::string& expires_on,
                                const std::optional<bool>& enabled,
                                std::string& error) {
    std::lock_guard<std::mutex> lock(mutex_);
    const auto account = FindAccount(database_, username);
    if (!account.has_value()) {
        error = "账号不存在";
        return false;
    }
    if (account->is_admin) {
        error = "内置管理员账号不能修改有效期";
        return false;
    }
    if (!expires_on.empty() && !IsValidDate(expires_on)) {
        error = "请输入有效日期";
        return false;
    }
    if (password.has_value() && !password->empty() &&
        !ValidatePassword(password.value(), error)) {
        return false;
    }
    std::string salt = account->password_salt;
    std::string password_hash = account->password_hash;
    const bool account_enabled = enabled.value_or(account->enabled);
    if (password.has_value() && !password->empty()) {
        salt = drogon::utils::secureRandomString(32);
        password_hash = HashPassword(password.value(), salt);
    }
    Statement statement(database_,
        "UPDATE accounts SET password_salt = ?, password_hash = ?, expires_on = ?, "
        "enabled = ? WHERE username = ?");
    if (!statement) {
        error = "保存账号失败";
        return false;
    }
    sqlite3_bind_text(statement.Get(), 1, salt.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 2, password_hash.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_text(statement.Get(), 3, expires_on.c_str(), -1, SQLITE_TRANSIENT);
    sqlite3_bind_int(statement.Get(), 4, account_enabled ? 1 : 0);
    sqlite3_bind_text(statement.Get(), 5, username.c_str(), -1, SQLITE_TRANSIENT);
    if (sqlite3_step(statement.Get()) != SQLITE_DONE) {
        error = "保存账号失败";
        return false;
    }
    InvalidateUserSessionsLocked(username);
    return true;
}

void AuthManager::InvalidateUserSessionsLocked(const std::string& username) {
    for (auto iterator = sessions_.begin(); iterator != sessions_.end();) {
        if (iterator->second.username == username) {
            iterator = sessions_.erase(iterator);
        } else {
            ++iterator;
        }
    }
}