Add enforcement for command invoke location

Equinox- · Feb 23, 2024 · 6cf1bf0 · 6cf1bf0
1 parent cf8c625
commit 6cf1bf0
Show file tree

Hide file tree

Showing 11 changed files with 68 additions and 19 deletions.
diff --git a/Meds.Watchdog/Discord/DiscordCmdBase.cs b/Meds.Watchdog/Discord/DiscordCmdBase.cs
@@ -0,0 +1,19 @@
+using System.Threading.Tasks;
+using DSharpPlus.SlashCommands;
+
+namespace Meds.Watchdog.Discord
+{
+    public abstract class DiscordCmdBase : ApplicationCommandModule
+    {
+        private readonly DiscordService _discord;
+
+        protected DiscordCmdBase(DiscordService discord)
+        {
+            _discord = discord;
+        }
+
+        public override Task<bool> BeforeSlashExecutionAsync(InteractionContext ctx) => _discord.VerifyRequirements(ctx);
+
+        public override Task<bool> BeforeContextMenuExecutionAsync(ContextMenuContext ctx) => _discord.VerifyRequirements(ctx);
+    }
+}
diff --git a/Meds.Watchdog/Discord/DiscordCmdDiagnostic.cs b/Meds.Watchdog/Discord/DiscordCmdDiagnostic.cs
@@ -8,12 +8,12 @@
 namespace Meds.Watchdog.Discord
 {
     // Not using command groups to allow for discrete permissions.
-    public class DiscordCmdDiagnostic : ApplicationCommandModule
+    public class DiscordCmdDiagnostic : DiscordCmdBase
     {
         private readonly DiagnosticController _diagnostic;
         private readonly ILogger<DiscordCmdDiagnostic> _log;
 
-        public DiscordCmdDiagnostic(DiagnosticController diagnostic, ILogger<DiscordCmdDiagnostic> log)
+        public DiscordCmdDiagnostic(DiagnosticController diagnostic, ILogger<DiscordCmdDiagnostic> log, DiscordService discord) : base(discord)
         {
             _diagnostic = diagnostic;
             _log = log;

diff --git a/Meds.Watchdog/Discord/DiscordCmdLifecycle.cs b/Meds.Watchdog/Discord/DiscordCmdLifecycle.cs
@@ -11,7 +11,7 @@
 namespace Meds.Watchdog.Discord
 {
     // Not using permission groups since they can't have discrete permissions
-    public class DiscordCmdLifecycle : ApplicationCommandModule
+    public class DiscordCmdLifecycle : DiscordCmdBase
     {
         private readonly LifecycleController _lifetimeController;
         private readonly HealthTracker _healthTracker;
@@ -20,7 +20,7 @@ public class DiscordCmdLifecycle : ApplicationCommandModule
         private readonly IHostApplicationLifetime _lifetime;
 
         public DiscordCmdLifecycle(LifecycleController lifetimeController, HealthTracker healthTracker,
-            InstallConfiguration configuration, ILogger<DiscordCmdLifecycle> log, IHostApplicationLifetime lifetime)
+            InstallConfiguration configuration, ILogger<DiscordCmdLifecycle> log, IHostApplicationLifetime lifetime, DiscordService discord) : base(discord)
         {
             _lifetimeController = lifetimeController;
             _healthTracker = healthTracker;

diff --git a/Meds.Watchdog/Discord/DiscordCmdPlayers.cs b/Meds.Watchdog/Discord/DiscordCmdPlayers.cs
@@ -8,15 +8,15 @@
 
 namespace Meds.Watchdog.Discord
 {
-    public class DiscordCmdPlayers : ApplicationCommandModule
+    public class DiscordCmdPlayers : DiscordCmdBase
     {
         private readonly ISubscriber<PlayersResponse> _playersSubscriber;
         private readonly IPublisher<PlayersRequest> _playersRequest;
         private readonly ISubscriber<PromotePlayerResponse> _promoteSubscriber;
         private readonly IPublisher<PromotePlayerRequest> _promoteRequest;
 
         public DiscordCmdPlayers(ISubscriber<PlayersResponse> playersSubscriber, IPublisher<PlayersRequest> playersRequest,
-            ISubscriber<PromotePlayerResponse> promoteSubscriber, IPublisher<PromotePlayerRequest> promoteRequest)
+            ISubscriber<PromotePlayerResponse> promoteSubscriber, IPublisher<PromotePlayerRequest> promoteRequest, DiscordService discord) : base(discord)
         {
             _playersSubscriber = playersSubscriber;
             _playersRequest = playersRequest;

diff --git a/Meds.Watchdog/Discord/DiscordCmdSave.cs b/Meds.Watchdog/Discord/DiscordCmdSave.cs
@@ -23,7 +23,7 @@
 namespace Meds.Watchdog.Discord
 {
     // Not using command groups for discrete permissions.
-    public class DiscordCmdSave : ApplicationCommandModule
+    public class DiscordCmdSave : DiscordCmdBase
     {
         private readonly ILogger<DiscordCmdSave> _log;
         private readonly HealthTracker _healthTracker;
@@ -38,7 +38,7 @@ public class DiscordCmdSave : ApplicationCommandModule
 
         public DiscordCmdSave(HealthTracker healthTracker, ISubscriber<SaveResponse> saveResponse, IPublisher<SaveRequest> saveRequest, SaveFiles saves,
             DataStore dataStore, ISubscriber<RestoreSceneResponse> restoreSceneResponse, IPublisher<RestoreSceneRequest> restoreSceneRequest,
-            LifecycleController lifecycle, ILogger<DiscordCmdSave> log, RtcFileSharing rtcFileSharing)
+            LifecycleController lifecycle, ILogger<DiscordCmdSave> log, RtcFileSharing rtcFileSharing, DiscordService discord) : base(discord)
         {
             _healthTracker = healthTracker;
             _saveResponse = saveResponse;

diff --git a/Meds.Watchdog/Discord/DiscordCmdSaveSearch.cs b/Meds.Watchdog/Discord/DiscordCmdSaveSearch.cs
@@ -14,13 +14,13 @@
 namespace Meds.Watchdog.Discord
 {
     // Not using command groups for discrete permissions.
-    public class DiscordCmdSaveSearch : ApplicationCommandModule
+    public class DiscordCmdSaveSearch : DiscordCmdBase
     {
         private readonly SaveFiles _saves;
         private readonly DataStore _dataStore;
 
         public DiscordCmdSaveSearch(SaveFiles saves,
-            DataStore dataStore)
+            DataStore dataStore, DiscordService discord) : base(discord)
         {
             _saves = saves;
             _dataStore = dataStore;

diff --git a/Meds.Watchdog/Discord/DiscordCmdStatus.cs b/Meds.Watchdog/Discord/DiscordCmdStatus.cs
@@ -9,12 +9,12 @@
 namespace Meds.Watchdog.Discord
 {
     // Not using command groups for discrete permissions
-    public class DiscordCmdStatus : ApplicationCommandModule
+    public class DiscordCmdStatus : DiscordCmdBase
     {
         private readonly LifecycleController _lifetimeController;
         private readonly HealthTracker _healthTracker;
 
-        public DiscordCmdStatus(LifecycleController lifetimeController, HealthTracker healthTracker)
+        public DiscordCmdStatus(LifecycleController lifetimeController, HealthTracker healthTracker, DiscordService discord) : base(discord)
         {
             _lifetimeController = lifetimeController;
             _healthTracker = healthTracker;

diff --git a/Meds.Watchdog/Discord/DiscordConfig.cs b/Meds.Watchdog/Discord/DiscordConfig.cs
@@ -14,6 +14,12 @@ public class DiscordConfig : MemberwiseEquatable<DiscordConfig>
 
         [XmlElement("ChannelSync")]
         public List<DiscordChannelSync> ChannelSyncs;
+
+        [XmlElement("RequireGuild")]
+        public List<ulong> RequireGuild;
+
+        [XmlElement("RequireChannel")]
+        public List<ulong> RequireChannel;
     }
 
     public class DiscordChannelSync : MemberwiseEquatable<DiscordChannelSync>

diff --git a/Meds.Watchdog/Discord/DiscordInit.cs b/Meds.Watchdog/Discord/DiscordInit.cs
@@ -1,4 +1,5 @@
 using System;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -21,6 +22,10 @@ public class DiscordService : IHostedService, IDisposable
         private readonly ILoggerFactory _rootLogger;
         private readonly IServiceProvider _services;
         private readonly ILogger<DiscordService> _log;
+        private readonly Refreshable<HashSet<ulong>> _requiredGuilds;
+        private readonly Refreshable<HashSet<ulong>> _requiredChannels;
+
+        private IDisposable _clientConnector;
 
         public DiscordClient Client => _state?.Client;
 
@@ -30,6 +35,8 @@ public DiscordService(Refreshable<Configuration> config, ILoggerFactory rootLogg
             _rootLogger = rootLogger;
             _services = provider;
             _log = log;
+            _requiredGuilds = _config.Map(x => new HashSet<ulong>(x.RequireGuild ?? new List<ulong>()));
+            _requiredChannels = _config.Map(x => new HashSet<ulong>(x.RequireChannel ?? new List<ulong>()));
         }
 
         private volatile State _state;
@@ -54,7 +61,6 @@ public State(DiscordService owner, string token)
                     ButtonBehavior = ButtonPaginationBehavior.DeleteButtons,
                     PaginationBehaviour = PaginationBehaviour.Ignore,
                     ResponseBehavior = InteractionResponseBehavior.Ack,
-                    AckPaginationButtons = true,
                 });
 
                 var commands = Client.UseSlashCommands(new SlashCommandsConfiguration
@@ -68,8 +74,7 @@ public State(DiscordService owner, string token)
                         args.Context.CommandName,
                         string.Join(", ", args.Context.Interaction?.Data?.Options?.Select(x => x.Name + " " + x.Value) ?? Array.Empty<string>()),
                         uuid);
-                    args.Context.FollowUpAsync(new DiscordFollowupMessageBuilder { Content = $"Command processing failed!  Error ID = {uuid}" });
-                    return Task.CompletedTask;
+                    return args.Context.FollowUpAsync(new DiscordFollowupMessageBuilder { Content = $"Command processing failed!  Error ID = {uuid}" });
                 };
 
                 commands.RegisterCommands<DiscordCmdDiagnostic>();
@@ -100,9 +105,6 @@ private async Task RefreshToken(Task prevTask, string token)
             await _state.Client.ConnectAsync();
         }
 
-
-        private IDisposable _clientConnector;
-
         public async Task StartAsync(CancellationToken cancellationToken)
         {
             _clientConnector = _config
@@ -124,5 +126,26 @@ public void Dispose()
         {
             Client?.Dispose();
         }
+
+        private bool CheckRequirements(BaseContext ctx)
+        {
+            var guilds = _requiredGuilds.Current;
+            var channels = _requiredChannels.Current;
+            if (guilds.Count == 0 && channels.Count == 0)
+                return true;
+            if (ctx.Guild != null && guilds.Contains(ctx.Guild.Id))
+                return true;
+            if (channels.Contains(ctx.Channel.Id))
+                return true;
+            return false;
+        }
+
+        internal async Task<bool> VerifyRequirements(BaseContext ctx)
+        {
+            if (CheckRequirements(ctx)) return true;
+            _log.ZLogWarning("Command {0} was used from an invalid location", ctx.CommandName);
+            await ctx.CreateResponseAsync("Command used from invalid location");
+            return false;
+        }
     }
 }
diff --git a/Meds.Watchdog/Discord/DiscordStatusMonitor.cs b/Meds.Watchdog/Discord/DiscordStatusMonitor.cs
@@ -31,7 +31,7 @@ public static string FormatStateRequest(LifecycleState request, bool readiness =
                 case LifecycleStateCase.Restarting:
                     return $"{request.Icon ?? "♻️"} | {request.Reason ?? "Restarting"}";
                 case LifecycleStateCase.Faulted:
-                    return $"{request.Icon ?? "🪦"} | {request.Reason ?? "Faulted"}";
+                    return $"{request.Icon ?? "💀"} | {request.Reason ?? "Faulted"}";
                 default:
                     throw new ArgumentOutOfRangeException();
             }

diff --git a/Meds.Watchdog/Meds.Watchdog.csproj b/Meds.Watchdog/Meds.Watchdog.csproj
@@ -48,6 +48,7 @@
         <Compile Include="Configuration.cs" />
         <Compile Include="DataStore.cs" />
         <Compile Include="DiagnosticController.cs" />
+        <Compile Include="Discord\DiscordCmdBase.cs" />
         <Compile Include="Discord\DiscordCmdDiagnostic.cs" />
         <Compile Include="Discord\DiscordCmdPlayers.cs" />
         <Compile Include="Discord\DiscordCmdSave.cs" />