mcdiscordbot/MinecraftDiscordBot/Program.cs

using CommandLine;
using Discord;
using Discord.Commands;
using Discord.Rest;
using Discord.Webhook;
using Discord.WebSocket;
using Fleck;
using MinecraftDiscordBot.Commands;
using MinecraftDiscordBot.Models;
using MinecraftDiscordBot.Services;
using System.Collections.Concurrent;
using System.Reflection;
using System.Runtime.CompilerServices;

namespace MinecraftDiscordBot;

public class Program : IDisposable, ICommandHandler<ResponseType>, IUserRoleManager {
    public const string WebSocketSource = "WebSocket";
    public const string BotSource = "Bot";
    private static readonly object LogLock = new();
    public const int ChoiceTimeout = 20 * 1000;
    private readonly DiscordSocketClient _client = new(new() {
        LogLevel = LogSeverity.Verbose,
        GatewayIntents = GatewayIntents.AllUnprivileged & ~(GatewayIntents.GuildScheduledEvents | GatewayIntents.GuildInvites)
    });
    private readonly WebSocketServer _wssv;
    private readonly BotConfiguration _config;
    private readonly HashSet<ulong> _whitelistedChannels;
    private readonly ConcurrentDictionary<Guid, RootCommandService> _connections = new();
    private static readonly char[] WhiteSpace = new char[] { '\t', '\n', ' ', '\r' };
    public IEnumerable<ActiveChannel> Channels => _channels ?? throw new InvalidProgramException("Channels used before verification!");
    public ActiveChannel[]? _channels;
    private bool disposedValue;
    private static ITextChannel? LogChannel;
    private readonly RootCommandService _computer;

    public static bool OnlineNotifications => true;
    public const LogSeverity DiscordLogSeverity = LogSeverity.Warning;
    private const string ClientScriptName = "MinecraftDiscordBot.ClientScript.lua";
    private const string WebhookName = "minecraftbot";
    public readonly string ClientScript;
    private readonly ITokenProvider _tokenProvider = new TimeoutTokenProvider(InstanceId, 10);
    private static readonly int InstanceId = new Random().Next();

    private string GetVerifiedClientScript() => ClientScript
        .Replace("$TOKEN", _tokenProvider.GenerateToken());

    private static string GetClientScript(BotConfiguration config) {
        using var stream = Assembly.GetExecutingAssembly().GetManifestResourceStream(ClientScriptName);
        if (stream is null) throw new FileNotFoundException("Client script could not be loaded!");
        using var sr = new StreamReader(stream);
        return sr.ReadToEnd()
            .Replace("$HOST", $"ws://{config.SocketHost}:{config.Port}");
    }

    private Task Broadcast(Func<ITextChannel, Task<IUserMessage>> message)
        => Task.WhenAll(Channels.Select(i => message(i.Channel)));
    public Program(BotConfiguration config) {
        _config = config;
        _computer = new(this);
        _computer.ChatMessageReceived += MinecraftMessageReceived;
        _computer.SocketChanged += ComputerConnectedChanged;
        _administrators = config.Administrators.ToHashSet();
        ClientScript = GetClientScript(config);
        _client.Log += LogAsync;
        _client.MessageReceived += (msg) => DiscordMessageReceived(msg, 20 * 1000);
        _client.ReactionAdded += DiscordReactionAdded;
        _wssv = new WebSocketServer($"ws://0.0.0.0:{config.Port}") {
            RestartAfterListenError = true
        };
        FleckLog.LogAction = LogWebSocket;
        _whitelistedChannels = config.Channels.ToHashSet();
    }

    private void ComputerConnectedChanged(object? sender, IWebSocketConnection? e)
        => _ = Task.Run(() => Broadcast(i => i.SendMessageAsync(e is not null
       ? "The Minecraft client is now available!"
       : "The Minecraft client disconnected!")));
    private void MinecraftMessageReceived(object? sender, ChatEvent e)
        => Task.Run(() => WebhookBroadcast(i => i.SendMessageAsync(e.Message, username: e.Username, avatarUrl: $"https://crafatar.com/renders/head/{e.UUID}")));
    private Task<T[]> WebhookBroadcast<T>(Func<DiscordWebhookClient, Task<T>> apply) => Task.WhenAll(Channels.Select(i => apply(new DiscordWebhookClient(i.Webhook))));
    private void LogWebSocket(LogLevel level, string message, Exception exception) => Log(new(level switch {
        LogLevel.Debug => LogSeverity.Debug,
        LogLevel.Info => LogSeverity.Info,
        LogLevel.Warn => LogSeverity.Warning,
        LogLevel.Error => LogSeverity.Error,
        _ => LogSeverity.Critical // Unknown logging states should behave critical
    }, WebSocketSource, message, exception));


    public static Task<int> Main(string[] args)
        => Parser.Default.ParseArguments<BotConfiguration, ConfigFile>(args)
        .MapResult<BotConfiguration, ConfigFile, Task<int>>(
            RunWithConfig,
            RunWithConfig,
            errs => Task.FromResult(1));

    private static Task<int> RunWithConfig(IBotConfigurator arg) => new Program(arg.Config).RunAsync();

    public async Task<int> RunAsync() {
        StartWebSocketServer();
        await _client.LoginAsync(TokenType.Bot, _config.Token);
        await _client.StartAsync();
        if (!await HasValidChannels())
            return 1;

        // Block this task until the program is closed.
        await Task.Delay(-1);
        return 0;
    }

    private async Task<bool> HasValidChannels() {
        if (_config.LogChannel is ulong logChannelId) {
            LogChannel = await IsValidChannel(logChannelId);
            if (LogChannel is null)
                await LogWarningAsync(BotSource, $"The given log channel ID is not valid '{logChannelId}'!");
        }
        if (await GetValidChannels(_whitelistedChannels) is not { Length: > 0 } channels) {
            await LogErrorAsync(BotSource, new InvalidOperationException("No valid textchannel was whitelisted!"));
            return false;
        }
        _channels = await Task.WhenAll(channels.Select(async i => new ActiveChannel(i, await GetOrCreateWebhook(i))));
        static async Task<IWebhook> GetOrCreateWebhook(ITextChannel i) {
            var hooks = (await i.GetWebhooksAsync()).Where(i => i.Name == WebhookName).FirstOrDefault();
            return hooks ?? await i.CreateWebhookAsync(WebhookName);
        }
        return true;
    }

    private void StartWebSocketServer() => _wssv.Start(socket => {
        socket.OnOpen = async () => await SocketOpened(socket);
        socket.OnClose = async () => await SocketClosed(socket);
        socket.OnMessage = async message => await SocketReceived(socket, message);
    });

    private async Task<ITextChannel[]> GetValidChannels(IEnumerable<ulong> ids)
        => (await Task.WhenAll(ids.Select(i => IsValidChannel(i)))).OfType<ITextChannel>().ToArray();
    private async Task<ITextChannel?> IsValidChannel(ulong channelId) {
        var channel = await _client.GetChannelAsync(channelId);
        if (channel is not ITextChannel textChannel) {
            if (channel is null) await LogWarningAsync(BotSource, $"Channel with id [{channelId}] does not exist!");
            else await LogWarningAsync(BotSource, $"Channel is not a text channels and will not be used: {channel.Name} [{channel.Id}]!");
            return null;
        }

        if (textChannel.Guild is RestGuild guild) {
            await guild.UpdateAsync();
            await LogInfoAsync(BotSource, $"Whitelisted in channel: {channel.Name} [{channel.Id}] on server {guild.Name} [{guild.Id}]");
        } else {
            await LogWarningAsync(BotSource, $"Whitelisted in channel: {channel.Name} [{channel.Id}] on unknown server!");
        }
        return textChannel;
    }

    private async Task SocketReceived(IWebSocketConnection socket, string message) {
        await LogInfoAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Received: {message}");
        await (message switch {
            "getcode" => SendClientCode(socket),
            string s when s.StartsWith("login=") => ClientComputerConnected(socket, s[6..]),
            string s when s.StartsWith("error=") => ClientComputerError(socket, s[6..]),
            _ => DisruptClientConnection(socket, "Protocol violation!")
        });
    }

    private static async Task ClientComputerError(IWebSocketConnection socket, string message)
        => await LogWarningAsync("Client", $"Computer failed to run the script: {message}");

    private async Task ClientComputerConnected(IWebSocketConnection socket, string token) {
        if (!_tokenProvider.VerifyToken(token)) {
            await DisruptClientConnection(socket, "outdated");
            return;
        }
        await LogInfoAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Client logged in with valid script!");
        AddComputerSocket(socket);
    }

    private static async Task DisruptClientConnection(IWebSocketConnection socket, string reason) {
        await socket.Send(reason);
        await LogWarningAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Client will be terminated, reason: {reason}");
        socket.Close();
    }

    private async Task SendClientCode(IWebSocketConnection socket) {
        await socket.Send(GetVerifiedClientScript());
        await LogInfoAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Script sent to client!");
    }

    private void AddComputerSocket(IWebSocketConnection socket) => _computer.Socket = socket;

    private void RemoveComputerSocket(IWebSocketConnection socket) {
        if (_computer.Socket is { ConnectionInfo.Id: Guid id } && id == socket.ConnectionInfo.Id) _computer.Socket = null;
    }

    private async Task SocketClosed(IWebSocketConnection socket) {
        RemoveComputerSocket(socket);
        await LogInfoAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Client disconnected!");
    }

    private static async Task SocketOpened(IWebSocketConnection socket) => await LogInfoAsync(WebSocketSource, $"[{socket.ConnectionInfo.Id}] Client connected from {socket.ConnectionInfo.ClientIpAddress}:{socket.ConnectionInfo.ClientPort}!");

    private async Task DiscordMessageReceived(SocketMessage arg, int timeout = 10000) {
        if (arg is not SocketUserMessage message) return;
        if (message.Author.IsBot) return;
        if (!IsChannelWhitelisted(arg.Channel)) return;
        if (arg.Type is not MessageType.Default) return;

        var cts = new CancellationTokenSource(timeout);

        if (IsCommand(message, out var argPos)) {
            await arg.Channel.TriggerTypingAsync();
            var parameters = message.Content[argPos..].Split(WhiteSpace, StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
            _ = Task.Run(async () => {
                var response = await HandleCommand(message, parameters, cts.Token);
                await SendResponse(message, response);
            });
            return;
        }

        await LogInfoAsync("Discord", $"[{arg.Author.Username}] {arg.Content}");
        _ = Task.Run(() => _computer.Chat.SendMessageAsync(arg.Content, arg.Author.Username, cts.Token));
    }

    private Task SendResponse(SocketUserMessage message, ResponseType response) => response switch {
        ResponseType.IChoiceResponse res => HandleChoice(message, res),
        ResponseType.StringResponse res => message.ReplyAsync(res.Message),
        ResponseType.FileResponse res => message.Channel.SendFileAsync(res.Path, text: res.Message),
        _ => message.ReplyAsync($"Whoops, someone forgot to implement '{response.GetType().Name}' responses?"),
    };

    private readonly ConcurrentDictionary<ulong, ResponseType.IChoiceResponse> _choiceWait = new();
    private readonly HashSet<ulong> _administrators;

    private async Task DiscordReactionAdded(Cacheable<IUserMessage, ulong> message, Cacheable<IMessageChannel, ulong> channel, SocketReaction reaction) {
        var msgObject = await message.GetOrDownloadAsync();
        if (reaction.UserId == _client.CurrentUser.Id) return;
        if (!_choiceWait.TryRemove(message.Id, out var choice)) { await LogInfoAsync(BotSource, "Reaction was added to message without choice object!"); return; }
        await msgObject.DeleteAsync();
        await LogInfoAsync(BotSource, $"Reaction {reaction.Emote.Name} was added to the choice by {reaction.UserId}!");
    }

    private async Task HandleChoice(SocketUserMessage message, ResponseType.IChoiceResponse res) {
        var reply = await message.ReplyAsync($"{res.Query}\n{string.Join("\n", res.Options)}");
        _choiceWait[reply.Id] = res;
        var reactions = new Emoji[] { new("0️⃣")/*, new("1️⃣"), new("2️⃣"), new("3️⃣"), new("4️⃣"), new("5️⃣"), new("6️⃣"), new("7️⃣"), new("8️⃣"), new("9️⃣")*/ };
        await reply.AddReactionsAsync(reactions);
        _ = Task.Run(async () => {
            await Task.Delay(ChoiceTimeout);
            _ = _choiceWait.TryRemove(message.Id, out _);
            await reply.ModifyAsync(i => i.Content = "You did not choose in time!");
            await reply.RemoveAllReactionsAsync();
        });
    }

    public async Task<ResponseType> HandleCommand(SocketUserMessage message, string[] parameters, CancellationToken ct) {
        if (_computer is ICommandHandler<ResponseType> handler)
            try {
                return await handler.HandleCommand(message, parameters, ct);
            } catch (TaskCanceledException) {
                return ResponseType.AsString("Your request could not be processed in time!");
            } catch (ReplyException e) {
                await LogInfoAsync(BotSource, e.Message);
                return ResponseType.AsString(e.Message);
            } catch (Exception e) {
                await LogErrorAsync(BotSource, e);
                return ResponseType.AsString($"Oopsie doopsie, this should not have happened!");
            }
        else return ResponseType.AsString("The Minecraft server is currently unavailable!");
    }

    private bool IsCommand(SocketUserMessage message, out int argPos) {
        argPos = 0;
        return message.HasStringPrefix(_config.Prefix, ref argPos);
    }
    private bool IsChannelWhitelisted(ISocketMessageChannel channel)
        => _whitelistedChannels.Contains(channel.Id);

    public static ConfiguredTaskAwaitable LogInfoAsync(string source, string message) => LogAsync(new(LogSeverity.Info, source, message)).ConfigureAwait(false);
    public static ConfiguredTaskAwaitable LogWarningAsync(string source, string message) => LogAsync(new(LogSeverity.Warning, source, message)).ConfigureAwait(false);
    public static ConfiguredTaskAwaitable LogErrorAsync(string source, Exception exception) => LogAsync(new(LogSeverity.Error, source, exception?.Message, exception)).ConfigureAwait(false);
    public static void LogInfo(string source, string message) => Log(new(LogSeverity.Info, source, message));
    public static void LogWarning(string source, string message) => Log(new(LogSeverity.Warning, source, message));
    public static void LogError(string source, Exception exception) => Log(new(LogSeverity.Error, source, exception?.Message, exception));

    private static async Task LogAsync(LogMessage msg) {
        lock (LogLock) {
            var oldColor = Console.ForegroundColor;
            try {
                Console.ForegroundColor = msg.Severity switch {
                    LogSeverity.Critical => ConsoleColor.Magenta,
                    LogSeverity.Error => ConsoleColor.Red,
                    LogSeverity.Warning => ConsoleColor.Yellow,
                    LogSeverity.Info => ConsoleColor.White,
                    LogSeverity.Verbose => ConsoleColor.Blue,
                    LogSeverity.Debug => ConsoleColor.DarkBlue,
                    _ => ConsoleColor.Cyan,
                };
                Console.WriteLine(msg.ToString());
            } finally {
                Console.ForegroundColor = oldColor;
            }
        }
        if (msg.Severity <= DiscordLogSeverity && LogChannel is ITextChannel log) {
            await log.SendMessageAsync($"{msg.Severity}: {msg}");
        }
    }

    public static void Log(LogMessage msg) => _ = Task.Run(() => LogAsync(msg));

    protected virtual void Dispose(bool disposing) {
        if (!disposedValue) {
            if (disposing) {
                // TODO: dispose managed state (managed objects)
                _wssv.Dispose();
                _client.Dispose();
            }

            // TODO: free unmanaged resources (unmanaged objects) and override finalizer
            // TODO: set large fields to null
            disposedValue = true;
        }
    }

    // // TODO: override finalizer only if 'Dispose(bool disposing)' has code to free unmanaged resources
    // ~Program()
    // {
    //     // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
    //     Dispose(disposing: false);
    // }

    public void Dispose() {
        // Do not change this code. Put cleanup code in 'Dispose(bool disposing)' method
        Dispose(disposing: true);
        GC.SuppressFinalize(this);
    }

    public void RequireAdministrator(ulong user, string? message = null) {
        if (!_administrators.Contains(user))
            throw new ReplyException(message ?? "User is not authorized to access this command!");
    }
}

public class ActiveChannel {
    public ActiveChannel(ITextChannel channel, IWebhook webhook) {
        Channel = channel;
        Webhook = webhook;
    }

    public IWebhook Webhook { get; }
    public ITextChannel Channel { get; }
}

public abstract class ResponseType {
    private static string DefaultDisplay<T>(T obj) => obj?.ToString() ?? throw new InvalidProgramException("ToString did not yield anything!");
    public static ResponseType AsString(string message) => new StringResponse(message);
    public static ResponseType FromChoice<T>(string query, IEnumerable<T> choice, Func<T, Task> resultHandler, Func<T, string>? display = null) => new ChoiceResponse<T>(query, choice, resultHandler, display ?? DefaultDisplay);
    internal static ResponseType File(string path, string message) => new FileResponse(path, message);

    public class StringResponse : ResponseType {
        public StringResponse(string message) => Message = message;
        public string Message { get; }
    }
    public interface IChoiceResponse {
        IEnumerable<string> Options { get; }
        string Query { get; }
        Task HandleResult(int index);
    }
    public class ChoiceResponse<T> : ResponseType, IChoiceResponse {
        private readonly Func<T, Task> _resultHandler;
        private readonly T[] _options;
        private readonly Func<T, string> _displayer;
        public IEnumerable<string> Options => _options.Select(_displayer);
        public string Query { get; }
        public Task HandleResult(int index) => _resultHandler(_options[index]);
        public ChoiceResponse(string query, IEnumerable<T> choice, Func<T, Task> resultHandler, Func<T, string> display) {
            Query = query;
            _resultHandler = resultHandler;
            _options = choice.ToArray();
            _displayer = display;
        }
    }

    public class FileResponse : ResponseType {
        public FileResponse(string path, string message) {
            Path = path;
            Message = message;
        }

        public string Path { get; }
        public string Message { get; }
    }
}