using System.Security.Cryptography; using System.Text; using Microsoft.Data.Sqlite; namespace ReverseLlama.Server; internal sealed class ManagementStore { private static readonly TimeSpan ApiKeyLastUsedWriteInterval = TimeSpan.FromMinutes(1); private readonly Dictionary _apiKeysByHash = new(StringComparer.Ordinal); private readonly string _connectionString = ""; private readonly string _databasePath = ""; private readonly object _lock = new(); private readonly ILogger _logger; private bool _isAvailable; private string? _lastError; public ManagementStore(ServerSettings settings, ILogger logger) { _logger = logger; try { _databasePath = ResolveDatabasePath(settings.ManagementDatabasePath); _connectionString = new SqliteConnectionStringBuilder { DataSource = _databasePath, Mode = SqliteOpenMode.ReadWriteCreate, Pooling = true }.ToString(); Initialize(); _isAvailable = true; } catch (Exception exception) { _lastError = exception.Message; _logger.LogError( exception, "Management database is disabled because SQLite could not be initialized at {DatabasePath}.", string.IsNullOrWhiteSpace(_databasePath) ? settings.ManagementDatabasePath : _databasePath); } } public string DatabasePath => _databasePath; public bool IsAvailable => _isAvailable; public string? LastError => _lastError; public bool HasApiKeys { get { if (!_isAvailable) { return false; } lock (_lock) { return _apiKeysByHash.Count > 0; } } } public bool IsApiKeyValid(string apiKey, bool updateLastUsed) { if (!_isAvailable || string.IsNullOrWhiteSpace(apiKey)) { return false; } var hash = HashApiKey(apiKey); var now = DateTimeOffset.UtcNow; lock (_lock) { if (!_apiKeysByHash.TryGetValue(hash, out var key)) { return false; } if (!updateLastUsed || key.LastUsedUtc is not null && now - key.LastUsedUtc.Value < ApiKeyLastUsedWriteInterval) { return true; } key.LastUsedUtc = now; try { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = "UPDATE api_keys SET last_used_at_utc = $last_used_at_utc WHERE id = $id"; command.Parameters.AddWithValue("$last_used_at_utc", now.ToString("O")); command.Parameters.AddWithValue("$id", key.Id); command.ExecuteNonQuery(); } catch (Exception exception) when (exception is SqliteException or IOException or UnauthorizedAccessException) { _logger.LogWarning(exception, "Failed to update API key last-used timestamp."); } return true; } } public IReadOnlyList ListApiKeys() { if (!_isAvailable) { return []; } lock (_lock) { return _apiKeysByHash.Values .OrderBy(key => key.Name, StringComparer.OrdinalIgnoreCase) .ThenBy(key => key.CreatedAtUtc) .Select(key => new ApiKeyInfo( key.Id, key.Name, key.KeyPrefix, key.CreatedAtUtc, key.LastUsedUtc)) .ToList(); } } public CreatedApiKey CreateApiKey(string? name) { EnsureAvailable(); var apiKey = GenerateApiKey(); var now = DateTimeOffset.UtcNow; var state = new ApiKeyState { Id = Guid.NewGuid().ToString("n"), Name = string.IsNullOrWhiteSpace(name) ? "API key" : name.Trim(), KeyHash = HashApiKey(apiKey), KeyPrefix = GetKeyPrefix(apiKey), CreatedAtUtc = now }; lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ INSERT INTO api_keys (id, name, key_hash, key_prefix, created_at_utc) VALUES ($id, $name, $key_hash, $key_prefix, $created_at_utc) """; command.Parameters.AddWithValue("$id", state.Id); command.Parameters.AddWithValue("$name", state.Name); command.Parameters.AddWithValue("$key_hash", state.KeyHash); command.Parameters.AddWithValue("$key_prefix", state.KeyPrefix); command.Parameters.AddWithValue("$created_at_utc", state.CreatedAtUtc.ToString("O")); command.ExecuteNonQuery(); _apiKeysByHash[state.KeyHash] = state; } return new CreatedApiKey( state.Id, state.Name, state.KeyPrefix, state.CreatedAtUtc, apiKey); } public bool DeleteApiKey(string id) { if (!_isAvailable || string.IsNullOrWhiteSpace(id)) { return false; } lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = "DELETE FROM api_keys WHERE id = $id"; command.Parameters.AddWithValue("$id", id); var deleted = command.ExecuteNonQuery() > 0; if (deleted) { foreach (var pair in _apiKeysByHash.Where(pair => pair.Value.Id == id).ToArray()) { _apiKeysByHash.Remove(pair.Key); } } return deleted; } } public ClientAccess GetClientAccess(string clientId) { if (!_isAvailable || string.IsNullOrWhiteSpace(clientId)) { return ClientAccess.Enabled; } lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ SELECT disabled_until_utc, disabled_manually, disabled_reason FROM client_controls WHERE client_id = $client_id """; command.Parameters.AddWithValue("$client_id", clientId); using var reader = command.ExecuteReader(); if (!reader.Read()) { return ClientAccess.Enabled; } var disabledUntil = ReadNullableDateTimeOffset(reader, 0); var disabledManually = reader.GetInt32(1) != 0; var reason = reader.IsDBNull(2) ? null : reader.GetString(2); if (disabledManually) { return new ClientAccess(true, null, true, reason); } if (disabledUntil is { } until && until > DateTimeOffset.UtcNow) { return new ClientAccess(true, until, false, reason); } return ClientAccess.Enabled; } } public IReadOnlyDictionary ListClientControls() { if (!_isAvailable) { return new Dictionary(StringComparer.OrdinalIgnoreCase); } var result = new Dictionary(StringComparer.OrdinalIgnoreCase); lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = "SELECT client_id, disabled_until_utc, disabled_manually, disabled_reason FROM client_controls"; using var reader = command.ExecuteReader(); while (reader.Read()) { var clientId = reader.GetString(0); var disabledUntil = ReadNullableDateTimeOffset(reader, 1); var disabledManually = reader.GetInt32(2) != 0; var reason = reader.IsDBNull(3) ? null : reader.GetString(3); result[clientId] = disabledManually ? new ClientAccess(true, null, true, reason) : disabledUntil is { } until && until > DateTimeOffset.UtcNow ? new ClientAccess(true, until, false, reason) : ClientAccess.Enabled; } } return result; } public void DisableClient(string clientId, TimeSpan? duration, bool manually, string? reason) { EnsureAvailable(); if (string.IsNullOrWhiteSpace(clientId)) { throw new ArgumentException("Client id is required.", nameof(clientId)); } var now = DateTimeOffset.UtcNow; var disabledUntil = manually ? null : now.Add(duration ?? TimeSpan.FromHours(1)).ToString("O"); lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ INSERT INTO client_controls ( client_id, disabled_until_utc, disabled_manually, disabled_reason, updated_at_utc) VALUES ( $client_id, $disabled_until_utc, $disabled_manually, $disabled_reason, $updated_at_utc) ON CONFLICT(client_id) DO UPDATE SET disabled_until_utc = excluded.disabled_until_utc, disabled_manually = excluded.disabled_manually, disabled_reason = excluded.disabled_reason, updated_at_utc = excluded.updated_at_utc """; command.Parameters.AddWithValue("$client_id", clientId); command.Parameters.AddWithValue("$disabled_until_utc", (object?)disabledUntil ?? DBNull.Value); command.Parameters.AddWithValue("$disabled_manually", manually ? 1 : 0); command.Parameters.AddWithValue("$disabled_reason", string.IsNullOrWhiteSpace(reason) ? DBNull.Value : reason.Trim()); command.Parameters.AddWithValue("$updated_at_utc", now.ToString("O")); command.ExecuteNonQuery(); } } public void EnableClient(string clientId) { EnsureAvailable(); if (string.IsNullOrWhiteSpace(clientId)) { throw new ArgumentException("Client id is required.", nameof(clientId)); } lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ INSERT INTO client_controls ( client_id, disabled_until_utc, disabled_manually, disabled_reason, updated_at_utc) VALUES ( $client_id, NULL, 0, NULL, $updated_at_utc) ON CONFLICT(client_id) DO UPDATE SET disabled_until_utc = NULL, disabled_manually = 0, disabled_reason = NULL, updated_at_utc = excluded.updated_at_utc """; command.Parameters.AddWithValue("$client_id", clientId); command.Parameters.AddWithValue("$updated_at_utc", DateTimeOffset.UtcNow.ToString("O")); command.ExecuteNonQuery(); } } public void RecordRequest(RequestMetric metric) { if (!_isAvailable) { return; } try { lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ INSERT INTO request_metrics ( client_id, model, method, path, status_code, token_count, started_at_utc, completed_at_utc, duration_ms) VALUES ( $client_id, $model, $method, $path, $status_code, $token_count, $started_at_utc, $completed_at_utc, $duration_ms) """; command.Parameters.AddWithValue("$client_id", metric.ClientId); command.Parameters.AddWithValue("$model", string.IsNullOrWhiteSpace(metric.Model) ? DBNull.Value : metric.Model); command.Parameters.AddWithValue("$method", metric.Method); command.Parameters.AddWithValue("$path", metric.Path); command.Parameters.AddWithValue("$status_code", metric.StatusCode is null ? DBNull.Value : metric.StatusCode.Value); command.Parameters.AddWithValue("$token_count", metric.TokenCount); command.Parameters.AddWithValue("$started_at_utc", metric.StartedAtUtc.ToString("O")); command.Parameters.AddWithValue("$completed_at_utc", metric.CompletedAtUtc.ToString("O")); command.Parameters.AddWithValue("$duration_ms", metric.Duration.TotalMilliseconds); command.ExecuteNonQuery(); } } catch (Exception exception) when (exception is SqliteException or IOException or UnauthorizedAccessException) { _logger.LogWarning(exception, "Failed to record request metric for client {ClientId}.", metric.ClientId); } } public IReadOnlyDictionary GetClientRequestStats() { if (!_isAvailable) { return new Dictionary(StringComparer.OrdinalIgnoreCase); } var now = DateTimeOffset.UtcNow; var since10 = now.AddMinutes(-10).ToString("O"); var sinceHour = now.AddHours(-1).ToString("O"); var result = new Dictionary(StringComparer.OrdinalIgnoreCase); lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ SELECT client_id, COUNT(*), SUM(CASE WHEN started_at_utc >= $since10 THEN 1 ELSE 0 END), SUM(CASE WHEN started_at_utc >= $sinceHour THEN 1 ELSE 0 END) FROM request_metrics GROUP BY client_id """; command.Parameters.AddWithValue("$since10", since10); command.Parameters.AddWithValue("$sinceHour", sinceHour); using var reader = command.ExecuteReader(); while (reader.Read()) { result[reader.GetString(0)] = new ClientRequestStats( reader.GetInt64(1), reader.GetInt64(2), reader.GetInt64(3)); } } return result; } public IReadOnlyDictionary GetModelUsageStats() { if (!_isAvailable) { return new Dictionary(StringComparer.OrdinalIgnoreCase); } var now = DateTimeOffset.UtcNow; var since10 = now.AddMinutes(-10).ToString("O"); var sinceHour = now.AddHours(-1).ToString("O"); var result = new Dictionary(StringComparer.OrdinalIgnoreCase); lock (_lock) { using var connection = OpenConnection(); using var command = connection.CreateCommand(); command.CommandText = """ SELECT model, COUNT(*), SUM(CASE WHEN started_at_utc >= $since10 THEN 1 ELSE 0 END), SUM(CASE WHEN started_at_utc >= $sinceHour THEN 1 ELSE 0 END), SUM(CASE WHEN started_at_utc >= $since10 THEN token_count ELSE 0 END), SUM(CASE WHEN started_at_utc >= $sinceHour THEN token_count ELSE 0 END) FROM request_metrics WHERE model IS NOT NULL AND model <> '' GROUP BY model """; command.Parameters.AddWithValue("$since10", since10); command.Parameters.AddWithValue("$sinceHour", sinceHour); using var reader = command.ExecuteReader(); while (reader.Read()) { result[reader.GetString(0)] = new ModelUsageStats( reader.GetInt64(1), reader.GetInt64(2), reader.GetInt64(3), reader.GetInt64(4), reader.GetInt64(5)); } } return result; } private void Initialize() { var directory = Path.GetDirectoryName(_databasePath); if (!string.IsNullOrWhiteSpace(directory)) { Directory.CreateDirectory(directory); } using var connection = OpenConnection(); using (var pragma = connection.CreateCommand()) { pragma.CommandText = "PRAGMA journal_mode=WAL"; pragma.ExecuteNonQuery(); } using (var command = connection.CreateCommand()) { command.CommandText = """ CREATE TABLE IF NOT EXISTS client_controls ( client_id TEXT NOT NULL PRIMARY KEY, disabled_until_utc TEXT NULL, disabled_manually INTEGER NOT NULL DEFAULT 0, disabled_reason TEXT NULL, updated_at_utc TEXT NOT NULL ); CREATE TABLE IF NOT EXISTS api_keys ( id TEXT NOT NULL PRIMARY KEY, name TEXT NOT NULL, key_hash TEXT NOT NULL UNIQUE, key_prefix TEXT NOT NULL, created_at_utc TEXT NOT NULL, last_used_at_utc TEXT NULL ); CREATE TABLE IF NOT EXISTS request_metrics ( id INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, client_id TEXT NOT NULL, model TEXT NULL, method TEXT NOT NULL, path TEXT NOT NULL, status_code INTEGER NULL, token_count INTEGER NOT NULL DEFAULT 0, started_at_utc TEXT NOT NULL, completed_at_utc TEXT NOT NULL, duration_ms REAL NOT NULL ); CREATE INDEX IF NOT EXISTS idx_request_metrics_client_started ON request_metrics (client_id, started_at_utc); CREATE INDEX IF NOT EXISTS idx_request_metrics_model_started ON request_metrics (model, started_at_utc); """; command.ExecuteNonQuery(); } using (var command = connection.CreateCommand()) { command.CommandText = """ SELECT id, name, key_hash, key_prefix, created_at_utc, last_used_at_utc FROM api_keys """; using var reader = command.ExecuteReader(); while (reader.Read()) { var state = new ApiKeyState { Id = reader.GetString(0), Name = reader.GetString(1), KeyHash = reader.GetString(2), KeyPrefix = reader.GetString(3), CreatedAtUtc = ReadDateTimeOffset(reader.GetString(4)), LastUsedUtc = ReadNullableDateTimeOffset(reader, 5) }; _apiKeysByHash[state.KeyHash] = state; } } _logger.LogInformation( "Loaded {ApiKeyCount} API key(s) from {DatabasePath}.", _apiKeysByHash.Count, _databasePath); } private SqliteConnection OpenConnection() { var connection = new SqliteConnection(_connectionString); connection.Open(); return connection; } private void EnsureAvailable() { if (!_isAvailable) { throw new InvalidOperationException(_lastError ?? "The management database is not available."); } } private static DateTimeOffset? ReadNullableDateTimeOffset(SqliteDataReader reader, int ordinal) => reader.IsDBNull(ordinal) ? null : ReadDateTimeOffset(reader.GetString(ordinal)); private static DateTimeOffset ReadDateTimeOffset(string value) => DateTimeOffset.TryParse(value, out var parsed) ? parsed : DateTimeOffset.MinValue; private static string GenerateApiKey() { var bytes = RandomNumberGenerator.GetBytes(32); return $"rl_{Base64UrlEncode(bytes)}"; } private static string HashApiKey(string apiKey) => Convert.ToHexString(SHA256.HashData(Encoding.UTF8.GetBytes(apiKey))).ToLowerInvariant(); private static string GetKeyPrefix(string apiKey) => apiKey.Length <= 12 ? apiKey : apiKey[..12]; private static string Base64UrlEncode(byte[] bytes) => Convert.ToBase64String(bytes) .TrimEnd('=') .Replace('+', '-') .Replace('/', '_'); private static string ResolveDatabasePath(string? configuredPath) { if (!string.IsNullOrWhiteSpace(configuredPath)) { var expanded = Environment.ExpandEnvironmentVariables(configuredPath); return Path.IsPathRooted(expanded) ? expanded : Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, expanded)); } return Path.Combine(AppContext.BaseDirectory, "App_Data", "management.sqlite"); } private sealed class ApiKeyState { public string Id { get; init; } = ""; public string Name { get; init; } = ""; public string KeyHash { get; init; } = ""; public string KeyPrefix { get; init; } = ""; public DateTimeOffset CreatedAtUtc { get; init; } public DateTimeOffset? LastUsedUtc { get; set; } } } internal sealed record ClientAccess( bool IsDisabled, DateTimeOffset? DisabledUntilUtc, bool DisabledManually, string? DisabledReason) { public static ClientAccess Enabled { get; } = new(false, null, false, null); } internal sealed record ApiKeyInfo( string Id, string Name, string KeyPrefix, DateTimeOffset CreatedAtUtc, DateTimeOffset? LastUsedUtc); internal sealed record CreatedApiKey( string Id, string Name, string KeyPrefix, DateTimeOffset CreatedAtUtc, string Key); internal sealed record RequestMetric( string ClientId, string? Model, string Method, string Path, int? StatusCode, int TokenCount, DateTimeOffset StartedAtUtc, DateTimeOffset CompletedAtUtc, TimeSpan Duration); internal sealed record ClientRequestStats( long Total, long Last10Minutes, long LastHour); internal sealed record ModelUsageStats( long TotalRequests, long RequestsLast10Minutes, long RequestsLastHour, long TokensLast10Minutes, long TokensLastHour);