noctbot/bot/src/plugins/violation.ts

import { Plugin, ICommandData, Command, Event, BotEventData } from "src/model/plugin";
import { parseArgs, parseDuration, UNIT_MEASURES, isAuthorisedAsync } from "src/util";
import { GuildMember, Guild, MessageEmbed, Message, TextChannel, PartialGuildMember, User } from "discord.js";
import { eventLogger, logger } from "src/logging";
import { client } from "src/client";
import humanizeDuration from "humanize-duration";
import { getRepository, ObjectType, FindConditions, DeepPartial } from "typeorm";
import { GuildViolationSettings } from "@shared/db/entity/GuildViolationSettings";
import { Mute, Violation } from "@shared/db/entity/Violation";
import { scheduleJob, Job, rescheduleJob } from "node-schedule";
import { QueryDeepPartialEntity } from "typeorm/query-builder/QueryPartialEntity";
import { tryDo, Option } from "@shared/common/async_utils";

const MENTION_PATTERN = /<@!?(\d+)>/;

interface ViolationInfo {
    member: GuildMember;
    endDate: Date;
    duration: number;
    guild: Guild;
    reason: string;
    settings: GuildViolationSettings;
    dryRun: boolean;
    noAnnounce: boolean;
}

type TimedViolation = Violation & { endsAt: Date };

type ModifyViolationFunction = (member: GuildMember | PartialGuildMember, settings: GuildViolationSettings, violation: DeepPartial<TimedViolation>) => DeepPartial<TimedViolation>; 
type StartViolationFunction = (member: GuildMember | PartialGuildMember, settings: GuildViolationSettings, violation: TimedViolation) => Promise<void>;
type StopViolationFunction = (guild: Guild, userId: string, settings: GuildViolationSettings, violation: TimedViolation) => Promise<void>;
interface TimedViolationStopHandler {
    type: ObjectType<TimedViolation>;
    start: StartViolationFunction;
    stop: StopViolationFunction;
    modify?: ModifyViolationFunction;
    command: string;
}

@Plugin
export class ViolationPlugin {
    jobs: Record<number, Job> = {};
    timedViolationHandlers: TimedViolationStopHandler[] = [
        {
            command: "mute",
            type: Mute,
            start: async (member: GuildMember | PartialGuildMember, settings: GuildViolationSettings, violation: TimedViolation): Promise<void> => {
                const muteRoleResolve = await tryDo(member.guild.roles.fetch(settings.muteRoleId));
                if (!muteRoleResolve.ok || !muteRoleResolve.result) {
                    logger.error(
                        "mute: Tried to mute user %s#%s (%s) but mute role ID %s is invalid!",
                        member.user?.username,
                        member.user?.discriminator,
                        member.user?.id,
                        settings.muteRoleId);
                    return;
                }
                // First mute, then remove other roles
                await member.roles.add(muteRoleResolve.result);
                const mute = violation as Mute;
                if (mute.previousRoles) {
                    const result = await tryDo(member.roles.remove(mute.previousRoles));
                    if (!result.ok) {
                        logger.error("mute: couldn't remove all roles from user %s#%s (%s)!" ,
                            member.user?.username,
                            member.user?.discriminator,
                            member.user?.id);
                    }
                }
            },
            stop: async (guild: Guild, userId: string, settings: GuildViolationSettings, violation: TimedViolation): Promise<void> => {
                const muteRoleResolve = await tryDo(guild.roles.fetch(settings.muteRoleId));

                if (!muteRoleResolve.ok || !muteRoleResolve.result) {
                    logger.warn("mute: couldn't find mute role id %s (removed from server?)", settings.muteRoleId);
                    return;
                }
                const muteRole = muteRoleResolve.result;

                const memberResolve = await tryDo(guild.members.fetch(userId));

                if (!memberResolve.ok) {
                    logger.warn("mute: user %s is not on the server anymore", userId);
                    return;
                }

                await memberResolve.result.roles.remove(muteRole);
                const mute = violation as Mute;
                if (mute.previousRoles) {
                    const result = await tryDo(memberResolve.result.roles.add(mute.previousRoles));
                    if (!result.ok) {
                        logger.warn("mute: couldn't readd all roles for user %s (tried to restore role ids: %s)", memberResolve.result.id, mute.previousRoles.join(", "));
                    }
                }
            },
            modify: (member: GuildMember | PartialGuildMember, settings: GuildViolationSettings, violation: DeepPartial<Mute>): DeepPartial<Mute> => {
                const originalRoles = member.roles.cache.keyArray().filter(r => r != settings.muteRoleId);
                violation.previousRoles = originalRoles;
                return violation;
            }
        }
    ];

    async start(): Promise<void> {
        for (const handler of this.timedViolationHandlers) {
            const repo = getRepository(handler.type);
            const validViolations = await repo.find({
                where: { valid: true }
            });

            for (const violation of validViolations) {
                const stopJob = this.scheduleRemoveViolation(handler.type, violation.guildId, violation.userId, handler.stop, handler.command);

                if (violation.endsAt <= new Date())
                    await stopJob();
                else
                    this.jobs[violation.id] = scheduleJob(violation.endsAt, stopJob);
            }
        }
    }

    @Event("guildMemberAdd")
    async onUserJoin(data: BotEventData, member: GuildMember | PartialGuildMember): Promise<void> {
        const settingsRepo = getRepository(GuildViolationSettings);
        const settings = await settingsRepo.findOne(member.guild.id);

        if (!settings)
            return;

        const hasActiveViolations = await getRepository(Violation).findOne({
            where: {
                guildId: member.guild.id,
                userId: member.id,
                valid: true
            }
        });

        if (!hasActiveViolations)
            return;

        for (const handler of this.timedViolationHandlers) {
            const repo = getRepository(handler.type);
            const activeViolations = await repo.find({
                where: {
                    guildId: member.guild.id,
                    userId: member.id,
                    valid: true
                }
            });

            if (activeViolations.length == 0)
                continue;

            for (const violation of activeViolations) {
                if (violation.endsAt < new Date())
                    await repo.update({ id: violation.id }, { valid: false });
                else
                    await handler.start(member, settings, violation);
            }
        }
    }

    @Command({
        type: "prefix",
        pattern: "mute",
        auth: true,
        documentation: {
            example: "mute[?!] <user> [<duration>] [<reason>]",
            description: "Mutes for a given duration and reason. ? = dry run, ! = no announcement"
        }
    })
    async muteUser({ message }: ICommandData): Promise<void> {
        await tryDo(message.delete());
        const info = await this.parseCommand(message, "mute");
        if (!info.ok)
            return;

        const handler = this.getViolationHandler(Mute);
        if (!handler) {
            logger.error("Couldn't find handler for Mute");
            return;
        }

        if (!info.dryRun) {
            eventLogger.warn("User %s#%s muted user %s#%s for %s because: %s", message.author.username, message.author.discriminator, info.member.user.username, info.member.user.discriminator, info.duration, info.reason);
        }
        await this.applyTimedViolation(Mute, info, "mute", handler.start, handler.stop, handler.modify);
        await this.sendViolationMessage(message, info, "User has been muted for server violation");
    }

    @Command({
        type: "prefix",
        pattern: "unmute",
        auth: true,
        documentation: {
            example: "unmute <user>",
            description: "Unmutes user"
        }
    })
    async unmuteUser({ message }: ICommandData): Promise<void> {
        await tryDo(message.delete());
        await this.removeTimedViolation(Mute, message, "mute");
    }

    private getViolationHandler(type: ObjectType<TimedViolation>): TimedViolationStopHandler {
        for (const handler of this.timedViolationHandlers) {
            if (handler.type == type)
                return handler;
        }
        throw new Error("Couldn't find handler for violation type!");
    }

    private async removeTimedViolation<T extends TimedViolation>(type: ObjectType<T>, message: Message, command = "violation") {
        if (!message.guild) {
            await message.reply("cannot do in DMs!");
            return;
        }

        const settingsRepo = getRepository(GuildViolationSettings);
        const settings = await settingsRepo.findOne(message.guild.id);
        if (!settings) {
            message.reply("this guild doesn't have violation settings set up!");
            return;
        }

        const [, userId] = parseArgs(message.content);
        if (!userId) {
            await message.reply("no user specified!");
            return;
        }
        if (userId == message.author.id) {
            await message.reply(`cannot ${command} yourself!`);
            return;
        }

        const user = await this.resolveUser(userId);
        if (!user) {
            await message.reply("couldn't find the given user!");
            logger.error("Tried to un-%s user %s but couldn't find them by id!", command, userId);
            return;
        }

        const violationRepo = getRepository(type);
        const existingViolation = await violationRepo.findOne({
            where: {
                guildId: message.guild.id,
                userId: user.id,
                valid: true
            }
        });
        if (!existingViolation) {
            await message.reply(`user has no existing active ${command}s in the DB!`);
            return;
        }

        await violationRepo.update({ id: existingViolation.id } as unknown as FindConditions<T>, { valid: false } as unknown as QueryDeepPartialEntity<T>);
        delete this.jobs[existingViolation.id];

        const handler = this.getViolationHandler(type);
        await handler.stop(message.guild, user.id, settings, existingViolation);
        await message.reply(`removed ${command} on user!`);
    }

    private async applyTimedViolation<T extends TimedViolation>(type: ObjectType<T>, info: ViolationInfo, command = "violation", apply: StartViolationFunction, remove: StopViolationFunction, modify?: ModifyViolationFunction) {
        if (info.dryRun)
            return;

        const violationRepo = getRepository(type);
        const existingViolation = await violationRepo.findOne({
            where: {
                userId: info.member.id,
                guildId: info.guild.id,
                valid: true
            }
        });

        let appliedViolation: T;
        if (existingViolation) {
            logger.warn("%s: trying to reapply on user %s#%s (%s)", command, info.member.user.username, info.member.user.discriminator, info.member.id);
            await violationRepo.update({ id: existingViolation.id } as unknown as FindConditions<T>, { endsAt: info.endDate } as unknown as QueryDeepPartialEntity<T>);
            const job = this.jobs[existingViolation.id];
            rescheduleJob(job, info.endDate);
            appliedViolation = existingViolation;
        } else {
            let rawViolation: DeepPartial<TimedViolation> = {
                guildId: info.guild.id,
                userId: info.member.id,
                reason: info.reason,
                endsAt: info.endDate,
                valid: true,
            };
            if (modify) {
                rawViolation = modify(info.member, info.settings, rawViolation);
            }
            const newViolation = await violationRepo.save(rawViolation as unknown as DeepPartial<T>);
            this.jobs[newViolation.id] = scheduleJob(info.endDate, this.scheduleRemoveViolation(type, info.guild.id, info.member.id, remove, command));
            appliedViolation = newViolation;
        }
        await apply(info.member, info.settings, appliedViolation);
    }

    private scheduleRemoveViolation<T extends TimedViolation>(type: ObjectType<T>, guildId: string, userId: string, handle: StopViolationFunction, command = "violation") {
        return async () => {
            const settingsRepo = getRepository(GuildViolationSettings);
            const settings = await settingsRepo.findOne(guildId);

            if (!settings) {
                logger.warn("un-%s: no violation settings found for guild %s", command, guildId);
                return;
            }

            const repo = getRepository(type);
            const violation = await repo.findOne({
                where: {
                    guildId: guildId,
                    userId: userId,
                    valid: true
                }
            });

            if (!violation) {
                logger.warn("un-%s: no violation found for user ID %s in guild %s", command, userId, guildId);
                return;
            }

            await repo.update({ id: violation.id } as unknown as FindConditions<T>, { valid: false } as unknown as QueryDeepPartialEntity<T>);
            delete this.jobs[violation.id];

            const guild = client.bot.guilds.resolve(guildId);
            if (!guild) {
                logger.warn("un-%s: couldn't find guild %s", command, guildId);
                return;
            }

            await handle(guild, userId, settings, violation);
        };
    }

    private async resolveUser(id: string): Promise<User | undefined> {
        const result = MENTION_PATTERN.exec(id);
        if (result) {
            const userId = result[1];
            const fetchResult = await tryDo(client.bot.users.fetch(userId));
            if (fetchResult.ok)
                return fetchResult.result;
        }

        const fetchResult = await tryDo(client.bot.users.fetch(id));
        if (!fetchResult.ok)
            return undefined;
        return fetchResult.result;
    }

    private async parseCommand(message: Message, command = "violation"): Promise<Option<ViolationInfo>> {
        if (!message.guild) {
            await message.reply("cannot do in DMs!");
            return { ok: false };
        }
        const violationSettingsRepo = getRepository(GuildViolationSettings);
        const settings = await violationSettingsRepo.findOne(message.guild.id);

        if (!settings) {
            await message.reply("sorry, this server doesn't have violation settings set up.");
            logger.error(
                "%s was called in guild %s (%s) on user %s which doesn't have config set up!",
                command,
                message.guild.name,
                message.guild.id,
                message.author.id);
            return { ok: false };
        }

        const [directive, userId, duration, ...rest] = parseArgs(message.content);

        const dryRun = directive.endsWith("?");
        const noAnnounce = directive.endsWith("!");

        if (!userId) {
            await message.reply("no user specified!");
            return { ok: false };
        }

        const user = await this.resolveUser(userId);

        if (!user) {
            await message.reply("couldn't find the given user!");
            logger.error("Tried to %s user %s but couldn't find them by id!", command, userId);
            return { ok: false };
        }

        if (user.id == message.author.id) {
            await message.reply(`cannot ${command} yourself!`);
            return { ok: false };
        }

        if (user.id == client.botUser.id) {
            await message.reply(`cannot apply ${command} on me!`);
            return { ok: false };
        }

        const memberResolve = await tryDo(message.guild.members.fetch(user));
        if (!memberResolve.ok) {
            await message.reply("user is not member of the server anymore!");
            logger.error("Tried to %s user %s but they are not on the server anymore!", command, userId);
            return { ok: false };
        }

        if (await isAuthorisedAsync(memberResolve.result)) {
            await message.reply(`cannot apply ${command} on another moderator!`);
            return { ok: false };
        }

        let durationMs = parseDuration(duration);
        let reasonArray = rest;

        if (!durationMs) {
            durationMs = UNIT_MEASURES.d as number;
            reasonArray = [duration, ...reasonArray];
        }

        const endDate = new Date(Date.now() + durationMs);
        let reason = reasonArray.join(" ");

        if (!reason)
            reason = "None given";

        return {
            ok: true,
            duration: durationMs,
            endDate: endDate,
            guild: message.guild,
            member: memberResolve.result,
            reason: reason,
            settings: settings,
            dryRun: dryRun,
            noAnnounce: noAnnounce
        };
    }

    private async sendViolationMessage(message: Message, info: ViolationInfo, title: string) {
        let announceChannel: TextChannel | null = null;
        if ((info.noAnnounce || info.dryRun) && message.channel.type == "text") {
            announceChannel = message.channel;
        }
        else if (info.settings.violationInfoChannelId) {
            const ch = info.guild.channels.resolve(info.settings.violationInfoChannelId);
            if (ch && ch.type == "text")
                announceChannel = ch as TextChannel;
            else if (message.channel.type == "text") {
                announceChannel = message.channel;
            }
        }
        await announceChannel?.send(new MessageEmbed({
            title: `${info.dryRun ? "[DRY RUN] " : ""}${title}`,
            color: 4944347,
            timestamp: new Date(),
            footer: {
                text: client.botUser.username
            },
            author: {
                name: client.botUser.username,
                iconURL: client.botUser.avatarURL() ?? undefined
            },
            fields: [
                {
                    name: "Username",
                    value: info.member.toString()
                },
                {
                    name: "Duration",
                    value: humanizeDuration(info.duration, { unitMeasures: UNIT_MEASURES })
                },
                {
                    name: "Reason",
                    value: info.reason
                }
            ]
        }));
    }
}