From 8e952672d87bd3fe414ba2ec9117cc96e5dab5d7 Mon Sep 17 00:00:00 2001 From: aNNiMON Date: Wed, 11 Jan 2023 22:20:12 +0200 Subject: [PATCH] Better permissions --- .../java/com/annimon/ffmpegbot/BotConfig.java | 5 ---- .../com/annimon/ffmpegbot/MainBotHandler.java | 11 +++----- .../com/annimon/ffmpegbot/Permissions.java | 25 +++++++++++++++++++ .../ffmpegbot/commands/admin/RunCommand.java | 3 ++- .../ffmpeg/MediaProcessingBundle.java | 9 ++++--- .../commands/ytdlp/YtDlpCommandBundle.java | 3 ++- 6 files changed, 38 insertions(+), 18 deletions(-) create mode 100644 src/main/java/com/annimon/ffmpegbot/Permissions.java diff --git a/src/main/java/com/annimon/ffmpegbot/BotConfig.java b/src/main/java/com/annimon/ffmpegbot/BotConfig.java index 76e0c8c..59d8d4b 100644 --- a/src/main/java/com/annimon/ffmpegbot/BotConfig.java +++ b/src/main/java/com/annimon/ffmpegbot/BotConfig.java @@ -4,9 +4,4 @@ import java.util.Set; public record BotConfig(String botToken, String botUsername, Set superUsers, Set allowedUsers) { - - boolean isUserAllowed(Long userId) { - return superUsers().contains(userId) - || allowedUsers().contains(userId); - } } diff --git a/src/main/java/com/annimon/ffmpegbot/MainBotHandler.java b/src/main/java/com/annimon/ffmpegbot/MainBotHandler.java index 7af793f..21b526d 100644 --- a/src/main/java/com/annimon/ffmpegbot/MainBotHandler.java +++ b/src/main/java/com/annimon/ffmpegbot/MainBotHandler.java @@ -19,11 +19,13 @@ import java.util.EnumSet; public class MainBotHandler extends BotHandler { private final BotConfig botConfig; + private final Permissions permissions; private final CommandRegistry commands; private final MediaProcessingBundle mediaProcessingBundle; public MainBotHandler(BotConfig botConfig) { this.botConfig = botConfig; + permissions = new Permissions(botConfig.superUsers(), botConfig.allowedUsers()); commands = new CommandRegistry<>(this, this::checkAccess); final var sessions = new Sessions(); mediaProcessingBundle = new MediaProcessingBundle(sessions); @@ -39,7 +41,7 @@ public class MainBotHandler extends BotHandler { if (commands.handleUpdate(update)) { return null; } - if (update.hasMessage()) { + if (update.hasMessage() && permissions.isUserAllowed(update.getMessage().getFrom().getId())) { mediaProcessingBundle.handleMessage(this, update.getMessage()); } return null; @@ -57,12 +59,7 @@ public class MainBotHandler extends BotHandler { private boolean checkAccess(Update update, @NotNull User user, @NotNull EnumSet roles) { final long userId = user.getId(); - - if (roles.contains(For.CREATOR) && botConfig.superUsers().contains(userId)) - return true; - if (roles.contains(For.ADMIN) && botConfig.isUserAllowed(userId)) - return true; - return roles.contains(For.USER); + return permissions.hasAccess(userId, roles); } @Override diff --git a/src/main/java/com/annimon/ffmpegbot/Permissions.java b/src/main/java/com/annimon/ffmpegbot/Permissions.java new file mode 100644 index 0000000..957570a --- /dev/null +++ b/src/main/java/com/annimon/ffmpegbot/Permissions.java @@ -0,0 +1,25 @@ +package com.annimon.ffmpegbot; + +import com.annimon.tgbotsmodule.commands.authority.For; +import org.jetbrains.annotations.NotNull; + +import java.util.EnumSet; +import java.util.Set; + +public record Permissions(Set superUsers, Set allowedUsers) { + public static final EnumSet SUPERUSERS = EnumSet.of(For.CREATOR); + public static final EnumSet ALLOWED_USERS = EnumSet.of(For.CREATOR, For.ADMIN); + + boolean hasAccess(long userId, @NotNull EnumSet roles) { + if (roles.contains(For.CREATOR) && superUsers().contains(userId)) + return true; + if (roles.contains(For.ADMIN) && isUserAllowed(userId)) + return true; + return roles.contains(For.USER); + } + + boolean isUserAllowed(Long userId) { + return superUsers().contains(userId) + || allowedUsers().contains(userId); + } +} diff --git a/src/main/java/com/annimon/ffmpegbot/commands/admin/RunCommand.java b/src/main/java/com/annimon/ffmpegbot/commands/admin/RunCommand.java index 83006a4..73ee598 100644 --- a/src/main/java/com/annimon/ffmpegbot/commands/admin/RunCommand.java +++ b/src/main/java/com/annimon/ffmpegbot/commands/admin/RunCommand.java @@ -1,5 +1,6 @@ package com.annimon.ffmpegbot.commands.admin; +import com.annimon.ffmpegbot.Permissions; import com.annimon.tgbotsmodule.commands.TextCommand; import com.annimon.tgbotsmodule.commands.authority.For; import com.annimon.tgbotsmodule.commands.context.MessageContext; @@ -25,7 +26,7 @@ public class RunCommand implements TextCommand { @SuppressWarnings("unchecked") @Override public EnumSet authority() { - return EnumSet.of(For.CREATOR); + return Permissions.SUPERUSERS; } @Override diff --git a/src/main/java/com/annimon/ffmpegbot/commands/ffmpeg/MediaProcessingBundle.java b/src/main/java/com/annimon/ffmpegbot/commands/ffmpeg/MediaProcessingBundle.java index 1c69a1d..56e6b2e 100644 --- a/src/main/java/com/annimon/ffmpegbot/commands/ffmpeg/MediaProcessingBundle.java +++ b/src/main/java/com/annimon/ffmpegbot/commands/ffmpeg/MediaProcessingBundle.java @@ -21,6 +21,7 @@ import java.util.function.BiConsumer; import java.util.function.Consumer; import java.util.stream.Collectors; +import static com.annimon.ffmpegbot.Permissions.ALLOWED_USERS; import static com.annimon.ffmpegbot.commands.ffmpeg.CallbackQueryCommands.*; import static com.annimon.ffmpegbot.commands.ffmpeg.MediaProcessingKeyboard.createKeyboard; @@ -34,10 +35,10 @@ public class MediaProcessingBundle implements CommandBundle { @Override public void register(@NotNull CommandRegistry commands) { - commands.register(new SimpleCallbackQueryCommand(PREV, For.ADMIN, ctx -> toggleParameter(ctx, true))); - commands.register(new SimpleCallbackQueryCommand(NEXT, For.ADMIN, ctx -> toggleParameter(ctx, false))); - commands.register(new SimpleCallbackQueryCommand(DETAIL, For.ADMIN, sessionCommand(this::details))); - commands.register(new SimpleCallbackQueryCommand(PROCESS, For.ADMIN, sessionCommand(this::process))); + commands.register(new SimpleCallbackQueryCommand(PREV, ALLOWED_USERS, ctx -> toggleParameter(ctx, true))); + commands.register(new SimpleCallbackQueryCommand(NEXT, ALLOWED_USERS, ctx -> toggleParameter(ctx, false))); + commands.register(new SimpleCallbackQueryCommand(DETAIL, ALLOWED_USERS, sessionCommand(this::details))); + commands.register(new SimpleCallbackQueryCommand(PROCESS, ALLOWED_USERS, sessionCommand(this::process))); } public void handleMessage(final @NotNull CommonAbsSender sender, final @NotNull Message message) { diff --git a/src/main/java/com/annimon/ffmpegbot/commands/ytdlp/YtDlpCommandBundle.java b/src/main/java/com/annimon/ffmpegbot/commands/ytdlp/YtDlpCommandBundle.java index 0514d41..c8499a5 100644 --- a/src/main/java/com/annimon/ffmpegbot/commands/ytdlp/YtDlpCommandBundle.java +++ b/src/main/java/com/annimon/ffmpegbot/commands/ytdlp/YtDlpCommandBundle.java @@ -1,5 +1,6 @@ package com.annimon.ffmpegbot.commands.ytdlp; +import com.annimon.ffmpegbot.Permissions; import com.annimon.ffmpegbot.session.*; import com.annimon.tgbotsmodule.commands.CommandBundle; import com.annimon.tgbotsmodule.commands.CommandRegistry; @@ -19,7 +20,7 @@ public class YtDlpCommandBundle implements CommandBundle { public void register(@NotNull CommandRegistry commands) { commands.register(new SimpleRegexCommand( Pattern.compile("/dl (https?://[^ ]+) ?(audio|\\d+)?p?"), - For.ADMIN, + Permissions.ALLOWED_USERS, this::download)); }