From a2f33b5df4e1cfe577df9bc3b515367024759a0a Mon Sep 17 00:00:00 2001 From: jvyden Date: Sat, 24 Aug 2024 16:24:29 -0400 Subject: [PATCH 1/4] AIPI integration --- .../Configuration/IntegrationConfig.cs | 23 ++- .../ApiTypes/Errors/ApiModerationError.cs | 10 ++ .../Endpoints/ApiV3/ResourceApiEndpoints.cs | 13 +- .../Endpoints/Game/PhotoEndpoints.cs | 14 +- Refresh.GameServer/RefreshContext.cs | 1 + Refresh.GameServer/RefreshGameServer.cs | 4 + Refresh.GameServer/Services/AipiService.cs | 161 ++++++++++++++++++ .../Services/DiscordStaffService.cs | 75 ++++++++ Refresh.GameServer/Services/ImportService.cs | 13 +- 9 files changed, 304 insertions(+), 10 deletions(-) create mode 100644 Refresh.GameServer/Endpoints/ApiV3/ApiTypes/Errors/ApiModerationError.cs create mode 100644 Refresh.GameServer/Services/AipiService.cs create mode 100644 Refresh.GameServer/Services/DiscordStaffService.cs diff --git a/Refresh.GameServer/Configuration/IntegrationConfig.cs b/Refresh.GameServer/Configuration/IntegrationConfig.cs index 56e4c382..1dcd6725 100644 --- a/Refresh.GameServer/Configuration/IntegrationConfig.cs +++ b/Refresh.GameServer/Configuration/IntegrationConfig.cs @@ -7,7 +7,7 @@ namespace Refresh.GameServer.Configuration; /// public class IntegrationConfig : Config { - public override int CurrentConfigVersion => 4; + public override int CurrentConfigVersion => 5; public override int Version { get; set; } protected override void Migrate(int oldVer, dynamic oldConfig) { @@ -31,11 +31,32 @@ protected override void Migrate(int oldVer, dynamic oldConfig) public bool DiscordWebhookEnabled { get; set; } public string DiscordWebhookUrl { get; set; } = "https://discord.com/api/webhooks/id/key"; + + public bool DiscordStaffWebhookEnabled { get; set; } + public string DiscordStaffWebhookUrl { get; set; } = "https://discord.com/api/webhooks/id/key"; public int DiscordWorkerFrequencySeconds { get; set; } = 60; public string DiscordNickname { get; set; } = "Refresh"; public string DiscordAvatarUrl { get; set; } = "https://raw.githubusercontent.com/LittleBigRefresh/Branding/main/icons/refresh_512x.png"; #endregion + #region AIPI + + public bool AipiEnabled { get; set; } = false; + public string AipiBaseUrl { get; set; } = "http://localhost:5000"; + + /// + /// The threshold at which tags are discarded during EVA2 prediction. + /// + public float AipiThreshold { get; set; } = 0.85f; + + // in DO we store this statically, but this exposing this as a config option allows us to obscure which tags + // are being blocked, because refresh is FOSS and DT could probably just look at it. + public string[] AipiBannedTags { get; set; } = []; + + public bool AipiRestrictAccountOnDetection { get; set; } = false; + + #endregion + public string? GrafanaDashboardUrl { get; set; } } \ No newline at end of file diff --git a/Refresh.GameServer/Endpoints/ApiV3/ApiTypes/Errors/ApiModerationError.cs b/Refresh.GameServer/Endpoints/ApiV3/ApiTypes/Errors/ApiModerationError.cs new file mode 100644 index 00000000..b5c07ef8 --- /dev/null +++ b/Refresh.GameServer/Endpoints/ApiV3/ApiTypes/Errors/ApiModerationError.cs @@ -0,0 +1,10 @@ +namespace Refresh.GameServer.Endpoints.ApiV3.ApiTypes.Errors; + +public class ApiModerationError : ApiError +{ + public static readonly ApiModerationError Instance = new(); + + public ApiModerationError() : base("This content was flagged as potentially unsafe, and administrators have been alerted. If you believe this is an error, please contact an administrator.", UnprocessableContent) + { + } +} \ No newline at end of file diff --git a/Refresh.GameServer/Endpoints/ApiV3/ResourceApiEndpoints.cs b/Refresh.GameServer/Endpoints/ApiV3/ResourceApiEndpoints.cs index c2c567eb..6e88834c 100644 --- a/Refresh.GameServer/Endpoints/ApiV3/ResourceApiEndpoints.cs +++ b/Refresh.GameServer/Endpoints/ApiV3/ResourceApiEndpoints.cs @@ -15,6 +15,7 @@ using Refresh.GameServer.Endpoints.ApiV3.DataTypes.Response; using Refresh.GameServer.Endpoints.ApiV3.DataTypes.Response.Data; using Refresh.GameServer.Importing; +using Refresh.GameServer.Services; using Refresh.GameServer.Types.Assets; using Refresh.GameServer.Types.Data; using Refresh.GameServer.Types.Roles; @@ -158,7 +159,11 @@ public ApiResponse UploadImageAsset(RequestContext context IDataStore dataStore, AssetImporter importer, GameServerConfig config, [DocSummary("The SHA1 hash of the asset")] string hash, - byte[] body, GameUser user, DataContext dataContext) + byte[] body, GameUser user, DataContext dataContext, + AipiService? aipi, + DiscordStaffService? discord, + IntegrationConfig integration + ) { // If we're blocking asset uploads, throw unless the user is an admin. // We also have the ability to block asset uploads for trusted users (when they would normally bypass this) @@ -197,6 +202,12 @@ public ApiResponse UploadImageAsset(RequestContext context return ApiInternalError.CouldNotWriteAssetError; gameAsset.OriginalUploader = user; + + if (aipi != null && aipi.ScanAndHandleAsset(dataContext, gameAsset)) + { + return ApiModerationError.Instance; + } + database.AddAssetToDatabase(gameAsset); return new ApiResponse(ApiGameAssetResponse.FromOld(gameAsset, dataContext)!, Created); diff --git a/Refresh.GameServer/Endpoints/Game/PhotoEndpoints.cs b/Refresh.GameServer/Endpoints/Game/PhotoEndpoints.cs index 92354099..61645767 100644 --- a/Refresh.GameServer/Endpoints/Game/PhotoEndpoints.cs +++ b/Refresh.GameServer/Endpoints/Game/PhotoEndpoints.cs @@ -6,6 +6,8 @@ using Bunkum.Protocols.Http; using Refresh.GameServer.Database; using Refresh.GameServer.Extensions; +using Refresh.GameServer.Services; +using Refresh.GameServer.Types.Assets; using Refresh.GameServer.Types.Data; using Refresh.GameServer.Types.Levels; using Refresh.GameServer.Types.Lists; @@ -19,7 +21,8 @@ public class PhotoEndpoints : EndpointGroup { [GameEndpoint("uploadPhoto", HttpMethods.Post, ContentType.Xml)] [RequireEmailVerified] - public Response UploadPhoto(RequestContext context, SerializedPhoto body, GameDatabaseContext database, GameUser user, IDataStore dataStore) + public Response UploadPhoto(RequestContext context, SerializedPhoto body, GameDatabaseContext database, GameUser user, IDataStore dataStore, + DataContext dataContext, AipiService aipi) { if (!dataStore.ExistsInStore(body.SmallHash) || !dataStore.ExistsInStore(body.MediumHash) || @@ -36,6 +39,15 @@ public Response UploadPhoto(RequestContext context, SerializedPhoto body, GameDa return BadRequest; } + List hashes = [body.LargeHash, body.MediumHash, body.SmallHash]; + foreach (string hash in hashes.Distinct()) + { + GameAsset? gameAsset = database.GetAssetFromHash(hash); + if(gameAsset == null) continue; + if (aipi != null && aipi.ScanAndHandleAsset(dataContext, gameAsset)) + return Unauthorized; + } + database.UploadPhoto(body, user); return OK; diff --git a/Refresh.GameServer/RefreshContext.cs b/Refresh.GameServer/RefreshContext.cs index 9c51a10f..3b905fba 100644 --- a/Refresh.GameServer/RefreshContext.cs +++ b/Refresh.GameServer/RefreshContext.cs @@ -9,4 +9,5 @@ public enum RefreshContext LevelListOverride, CoolLevels, Publishing, + Aipi, } \ No newline at end of file diff --git a/Refresh.GameServer/RefreshGameServer.cs b/Refresh.GameServer/RefreshGameServer.cs index 879fe910..fb7a5f81 100644 --- a/Refresh.GameServer/RefreshGameServer.cs +++ b/Refresh.GameServer/RefreshGameServer.cs @@ -139,6 +139,10 @@ protected override void SetupServices() this.Server.AddService(); this.Server.AddService(); this.Server.AddService(); + this.Server.AddService(); + + if(this._integrationConfig!.AipiEnabled) + this.Server.AddService(); #if DEBUG this.Server.AddService(); diff --git a/Refresh.GameServer/Services/AipiService.cs b/Refresh.GameServer/Services/AipiService.cs new file mode 100644 index 00000000..95357e84 --- /dev/null +++ b/Refresh.GameServer/Services/AipiService.cs @@ -0,0 +1,161 @@ +using System.Diagnostics; +using System.Net.Http.Json; +using Bunkum.Core.Services; +using JetBrains.Annotations; +using NotEnoughLogs; +using Refresh.GameServer.Configuration; +using Refresh.GameServer.Importing; +using Refresh.GameServer.Types.Assets; +using Refresh.GameServer.Types.Data; +using SixLabors.ImageSharp; +using SixLabors.ImageSharp.Formats; +using SixLabors.ImageSharp.Processing; + +namespace Refresh.GameServer.Services; + +// Referenced from DO. +public class AipiService : EndpointService +{ + private readonly HttpClient _client; + private readonly IntegrationConfig _config; + private readonly DiscordStaffService? _discord; + + private readonly ImageImporter _importer; + + [UsedImplicitly] + public AipiService(Logger logger, IntegrationConfig config, ImportService import, DiscordStaffService discord) : base(logger) + { + this._discord = discord; + this._config = config; + + this._client = new HttpClient + { + BaseAddress = new Uri(config.AipiBaseUrl), + }; + + this._importer = import.ImageImporter; + } + + public override void Initialize() + { + if (!this._config.DiscordStaffWebhookEnabled) + { + this.Logger.LogWarning(RefreshContext.Aipi, + "The Discord staff webhook is not enabled, but AIPI is. This is probably behavior you don't want."); + } + this.TestConnectivityAsync().Wait(); + } + + private async Task TestConnectivityAsync() + { + try + { + HttpResponseMessage response = await this._client.GetAsync("/"); + string content = await response.Content.ReadAsStringAsync(); + + if (response.IsSuccessStatusCode && content == "AIPI scanning service") + this.Logger.LogInfo(RefreshContext.Aipi, "AIPI appears to be working correctly"); + else + this.Logger.LogError(RefreshContext.Aipi, + $"AIPI seems to be down. Status code: {response.StatusCode}, content: {content}"); + } + catch (Exception e) + { + this.Logger.LogError(RefreshContext.Aipi, "AIPI connection failed: {0}", e.ToString()); + } + } + + private async Task PostAsync(string endpoint, Stream data) + { + HttpResponseMessage response = await this._client.PostAsync(endpoint, new StreamContent(data)); + AipiResponse? aipiResponse = await response.Content.ReadFromJsonAsync>(); + + if (aipiResponse == null) throw new Exception("No response was received from the server."); + if (!aipiResponse.Success) throw new Exception($"{response.StatusCode}: {aipiResponse.Reason}"); + + return aipiResponse.Data!; + } + + private async Task> PredictEvaAsync(Stream data) + { + Stopwatch stopwatch = new(); + this.Logger.LogTrace(RefreshContext.Aipi, "Pre-processing image data..."); + + DecoderOptions options = new() + { + MaxFrames = 1, + Configuration = SixLabors.ImageSharp.Configuration.Default, + }; + + Image image = await Image.LoadAsync(options, data); + // Technically, we don't read videos in Refresh like in DO, but a couple of users are currently using APNGs as their avatar. + // I don't want to break APNGs as they're harmless, so let's handle this by just reading the first frame for now. + if (image.Frames.Count > 0) + image = image.Frames.CloneFrame(0); + + image.Mutate(x => x.Resize(new ResizeOptions + { + Size = new Size(512), + Mode = ResizeMode.Max, + })); + + using MemoryStream processedData = new(); + await image.SaveAsPngAsync(processedData); + // await image.SaveAsPngAsync($"/tmp/{DateTimeOffset.Now.ToUnixTimeMilliseconds()}.png"); + processedData.Seek(0, SeekOrigin.Begin); + + float threshold = this._config.AipiThreshold; + + this.Logger.LogDebug(RefreshContext.Aipi, $"Running prediction for image @ threshold={threshold}..."); + + stopwatch.Start(); + Dictionary prediction = await this.PostAsync>($"/eva/predict?threshold={threshold}", processedData); + stopwatch.Stop(); + + this.Logger.LogInfo(RefreshContext.Aipi, $"Got prediction result in {stopwatch.ElapsedMilliseconds}ms."); + this.Logger.LogDebug(RefreshContext.Aipi, JsonConvert.SerializeObject(prediction)); + return prediction; + } + + public bool ScanAndHandleAsset(DataContext context, GameAsset asset) + { + // guard the fact that assets have an owner + Debug.Assert(asset.OriginalUploader != null, $"Asset {asset.AssetHash} had no original uploader when trying to scan"); + if (asset.OriginalUploader == null) + return false; + + // import the asset as png + bool isPspAsset = asset.AssetHash.StartsWith("psp/"); + string hash = isPspAsset ? "psp/" + asset.AssetHash : asset.AssetHash; + + if (!context.DataStore.ExistsInStore("png/" + asset.AssetHash)) + { + this._importer.ImportAsset(asset.AssetHash, isPspAsset, asset.AssetType, context.DataStore); + } + + // do actual prediction + using Stream stream = context.DataStore.GetStreamFromStore("png/" + asset.AssetHash); + Dictionary results = this.PredictEvaAsync(stream).Result; + + if (!results.Any(r => this._config.AipiBannedTags.Contains(r.Key))) + return false; + + this._discord?.PostPredictionResult(results, asset); + + if (this._config.AipiRestrictAccountOnDetection) + { + const string reason = "Automatic restriction for posting disallowed content. This will usually be undone within 24 hours if this is a mistake."; + context.Database.RestrictUser(asset.OriginalUploader, reason, DateTimeOffset.MaxValue); + } + + return true; + } + + private class AipiResponse + { + public bool Success { get; set; } + + public TData? Data { get; set; } + public string? Reason { get; set; } + } +} \ No newline at end of file diff --git a/Refresh.GameServer/Services/DiscordStaffService.cs b/Refresh.GameServer/Services/DiscordStaffService.cs new file mode 100644 index 00000000..16a47e59 --- /dev/null +++ b/Refresh.GameServer/Services/DiscordStaffService.cs @@ -0,0 +1,75 @@ +using System.Diagnostics; +using Bunkum.Core.Services; +using Discord; +using Discord.Webhook; +using NotEnoughLogs; +using Refresh.GameServer.Configuration; +using Refresh.GameServer.Types.UserData; +using GameAsset = Refresh.GameServer.Types.Assets.GameAsset; + +namespace Refresh.GameServer.Services; + +public class DiscordStaffService : EndpointService +{ + private readonly DiscordWebhookClient? _client; + private readonly IntegrationConfig _config; + + private readonly string _externalUrl; + + private const string NameSuffix = " (Staff)"; + + private const string DefaultResultsDescription = "These are the results of the AI's best guess at deciphering the contents of the image. " + + "Take them with a grain of salt as the AI isn't perfect."; + + internal DiscordStaffService(Logger logger, GameServerConfig gameConfig, IntegrationConfig config) : base(logger) + { + this._config = config; + this._externalUrl = gameConfig.WebExternalUrl; + + if(config.DiscordStaffWebhookEnabled) + this._client = new DiscordWebhookClient(config.DiscordStaffWebhookUrl); + } + + private string GetAssetUrl(string hash) + { + return $"{this._externalUrl}/api/v3/assets/{hash}/image"; + } + + private string GetAssetInfoUrl(string hash) + { + return $"{this._externalUrl}/api/v3/assets/{hash}/image"; + } + + private void PostMessage(string? message = null, IEnumerable? embeds = null!) + { + if (this._client == null) + return; + + embeds ??= []; + + ulong id = this._client.SendMessageAsync(embeds: embeds, + username: this._config.DiscordNickname + NameSuffix, avatarUrl: this._config.DiscordAvatarUrl).Result; + + this.Logger.LogInfo(RefreshContext.Discord, $"Posted webhook {id}: '{message}'"); + } + + public void PostPredictionResult(Dictionary results, GameAsset asset) + { + GameUser author = asset.OriginalUploader!; + + EmbedBuilder builder = new EmbedBuilder() + .WithAuthor($"Image posted by @{author.Username} (id: {author.UserId})", this.GetAssetUrl(author.IconHash)) + .WithDescription(DefaultResultsDescription) + .WithUrl(this.GetAssetInfoUrl(asset.AssetHash)) + .WithTitle($"AI Analysis of `{asset.AssetHash}`"); + + foreach ((string tag, float confidence) in results.OrderByDescending(r => r.Value).Take(25)) + { + string tagFormatted = this._config.AipiBannedTags.Contains(tag) ? $"{tag} (flagged!)" : tag; + string confidenceFormatted = confidence.ToString("0.00%"); + builder.AddField(tagFormatted, confidenceFormatted, true); + } + + this.PostMessage($"Prediction result for {author.Username}:", [builder.Build()]); + } +} \ No newline at end of file diff --git a/Refresh.GameServer/Services/ImportService.cs b/Refresh.GameServer/Services/ImportService.cs index 57e2b85e..71f7b32e 100644 --- a/Refresh.GameServer/Services/ImportService.cs +++ b/Refresh.GameServer/Services/ImportService.cs @@ -1,4 +1,3 @@ -using System.Reflection; using Bunkum.Core; using Bunkum.Listener.Request; using Bunkum.Core.Database; @@ -12,20 +11,20 @@ public class ImportService : Service { internal ImportService(Logger logger, TimeProviderService timeProvider) : base(logger) { - this._assetImporter = new AssetImporter(logger, timeProvider.TimeProvider); - this._imageImporter = new ImageImporter(logger); + this.AssetImporter = new AssetImporter(logger, timeProvider.TimeProvider); + this.ImageImporter = new ImageImporter(logger); } - private readonly AssetImporter _assetImporter; - private readonly ImageImporter _imageImporter; + internal readonly AssetImporter AssetImporter; + internal readonly ImageImporter ImageImporter; public override object? AddParameterToEndpoint(ListenerContext context, BunkumParameterInfo paramInfo, Lazy database) { if (paramInfo.ParameterType == typeof(AssetImporter)) - return this._assetImporter; + return this.AssetImporter; if (paramInfo.ParameterType == typeof(ImageImporter)) - return this._imageImporter; + return this.ImageImporter; return base.AddParameterToEndpoint(context, paramInfo, database); } From 36dbfc51d84f4bc73aedef57a836c4d37707cf07 Mon Sep 17 00:00:00 2001 From: jvyden Date: Sat, 24 Aug 2024 18:28:15 -0400 Subject: [PATCH 2/4] Update Refresh.GameServer/Services/DiscordStaffService.cs Co-authored-by: Beyley Thomas Signed-off-by: jvyden --- Refresh.GameServer/Services/DiscordStaffService.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Refresh.GameServer/Services/DiscordStaffService.cs b/Refresh.GameServer/Services/DiscordStaffService.cs index 16a47e59..a3bea12d 100644 --- a/Refresh.GameServer/Services/DiscordStaffService.cs +++ b/Refresh.GameServer/Services/DiscordStaffService.cs @@ -70,6 +70,6 @@ public void PostPredictionResult(Dictionary results, GameAsset as builder.AddField(tagFormatted, confidenceFormatted, true); } - this.PostMessage($"Prediction result for {author.Username}:", [builder.Build()]); + this.PostMessage($"Prediction result for {asset.AssetHash} ({author.Username}):", [builder.Build()]); } } \ No newline at end of file From e21044908054ebc14d5c597532afef148f580ec8 Mon Sep 17 00:00:00 2001 From: jvyden Date: Sat, 24 Aug 2024 18:28:40 -0400 Subject: [PATCH 3/4] Update Refresh.GameServer/Services/DiscordStaffService.cs Co-authored-by: Beyley Thomas Signed-off-by: jvyden --- Refresh.GameServer/Services/DiscordStaffService.cs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Refresh.GameServer/Services/DiscordStaffService.cs b/Refresh.GameServer/Services/DiscordStaffService.cs index a3bea12d..faadebc1 100644 --- a/Refresh.GameServer/Services/DiscordStaffService.cs +++ b/Refresh.GameServer/Services/DiscordStaffService.cs @@ -37,7 +37,7 @@ private string GetAssetUrl(string hash) private string GetAssetInfoUrl(string hash) { - return $"{this._externalUrl}/api/v3/assets/{hash}/image"; + return $"{this._externalUrl}/api/v3/assets/{hash}"; } private void PostMessage(string? message = null, IEnumerable? embeds = null!) From 5fd4f26bae7637bcd8fc576a69590f25d345c1aa Mon Sep 17 00:00:00 2001 From: jvyden Date: Sat, 24 Aug 2024 18:29:04 -0400 Subject: [PATCH 4/4] Remove unused variable --- Refresh.GameServer/Services/AipiService.cs | 1 - 1 file changed, 1 deletion(-) diff --git a/Refresh.GameServer/Services/AipiService.cs b/Refresh.GameServer/Services/AipiService.cs index 95357e84..9152f85c 100644 --- a/Refresh.GameServer/Services/AipiService.cs +++ b/Refresh.GameServer/Services/AipiService.cs @@ -126,7 +126,6 @@ public bool ScanAndHandleAsset(DataContext context, GameAsset asset) // import the asset as png bool isPspAsset = asset.AssetHash.StartsWith("psp/"); - string hash = isPspAsset ? "psp/" + asset.AssetHash : asset.AssetHash; if (!context.DataStore.ExistsInStore("png/" + asset.AssetHash)) {