matrix-spam-ml

git clone git://archive.git.mtrnord.blog/MTRNord/matrix-spam-ml.git
Log | Files | Refs | Submodules | README | LICENSE

index.ts (11442B)


      1 import {
      2     MatrixClient,
      3     SimpleFsStorageProvider,
      4     AutojoinRoomsMixin,
      5     RustSdkCryptoStorageProvider,
      6     LogService,
      7     RichConsoleLogger,
      8     MessageEvent,
      9     MessageEventContent,
     10     MatrixProfileInfo,
     11     CanonicalAliasEventContent,
     12     RoomNameEventContent,
     13 } from "matrix-bot-sdk";
     14 import { readFile } from "fs/promises";
     15 import { load } from "js-yaml";
     16 import * as tf from "@tensorflow/tfjs-node";
     17 import { Rank, Tensor } from "@tensorflow/tfjs-node";
     18 import { TFSavedModel } from "@tensorflow/tfjs-node/dist/saved_model";
     19 import { ReactionEvent } from "./events/ReactionEvent";
     20 import { htmlToText } from "html-to-text";
     21 import { BAN_REACTION, FALSE_POSITIVE_REACTION, getReactionHandler, KICK_REACTION } from "./reactionHandlers/reactionHandler";
     22 import BanListHandler from "./banlist/banlist";
     23 
     24 export type Config = {
     25     homeserver: string;
     26     accessToken: string;
     27     modelPath: string;
     28     // A private room for admins to see responses on reports and other secret activity.
     29     adminRoom: string;
     30     // A possibly public room where warnings land which also is used to issue actions from as an admin
     31     warningsRoom: string;
     32     // TODO: Dont set this via config but via a setup in the room on first launch or via the cli or something.
     33     banlistRoom: string;
     34 };
     35 
     36 const THRESHOLD = 0.8;
     37 const startUpMessage = "Bot is starting up...";
     38 
     39 class Bot {
     40     private readonly policyRoomHandler = new BanListHandler(this.client, this.config);
     41     public static async createBot() {
     42         LogService.setLogger(new RichConsoleLogger());
     43         LogService.muteModule("Metrics");
     44         const config = load(await readFile("./config.yaml", "utf8")) as Config;
     45 
     46         // Add some divider for clarity after tensorflow loaded up
     47         const line = '-'.repeat(process.stdout.columns);
     48         console.log(line);
     49 
     50         // Check if required fields are set
     51         if (config.homeserver == undefined && config.accessToken == undefined) {
     52             LogService.error("index", "Missing homeserver and accessToken config values");
     53             process.exit(1);
     54         } else if (config.homeserver == undefined) {
     55             LogService.error("index", "Missing homeserver config value");
     56             process.exit(1);
     57         } else if (config.accessToken == undefined) {
     58             LogService.error("index", "Missing accessToken config value");
     59             process.exit(1);
     60         }
     61 
     62         if (config.adminRoom == undefined && config.warningsRoom == undefined) {
     63             LogService.error("index", "Missing adminRoom and warningsRoom config values");
     64             process.exit(1);
     65         } else if (config.adminRoom == undefined) {
     66             LogService.error("index", "Missing adminRoom config value");
     67             process.exit(1);
     68         } else if (config.warningsRoom == undefined) {
     69             LogService.error("index", "Missing warningsRoom config value");
     70             process.exit(1);
     71         }
     72 
     73         if (config.adminRoom.startsWith("#")) {
     74             LogService.error("index", "adminRoom config value needs to be a roomid starting with a \"!\"");
     75             process.exit(1);
     76         }
     77         if (config.warningsRoom.startsWith("#")) {
     78             LogService.error("index", "warningsRoom config value needs to be a roomid starting with a \"!\"");
     79             process.exit(1);
     80         }
     81 
     82 
     83         // This will be the URL where clients can reach your homeserver. Note that this might be different
     84         // from where the web/chat interface is hosted. The server must support password registration without
     85         // captcha or terms of service (public servers typically won't work).
     86         const homeserverUrl = config.homeserver;
     87 
     88         // Use the access token you got from login or registration above.
     89         const accessToken = config.accessToken;
     90 
     91         // In order to make sure the bot doesn't lose its state between restarts, we'll give it a place to cache
     92         // any information it needs to. You can implement your own storage provider if you like, but a JSON file
     93         // will work fine for this example.
     94         const storage = new SimpleFsStorageProvider("ml-bot.json");
     95         const cryptoProvider = new RustSdkCryptoStorageProvider("./ml-bot-store");
     96 
     97         tf.enableProdMode()
     98         const model = await tf.node.loadSavedModel(config.modelPath);
     99 
    100         // Finally, let's create the client and set it to autojoin rooms. Autojoining is typical of bots to ensure
    101         // they can be easily added to any room.
    102         const client = new MatrixClient(homeserverUrl, accessToken, storage, cryptoProvider);
    103 
    104         // TODO: replace with manual handling to need admin approval
    105         AutojoinRoomsMixin.setupOnClient(client);
    106 
    107         // Join rooms as needed but crash if missing
    108         await client.joinRoom(config.adminRoom);
    109         await client.joinRoom(config.warningsRoom);
    110 
    111 
    112         return new Bot(config, client, model);
    113     }
    114 
    115     private constructor(private config: Config, private client: MatrixClient, private model: TFSavedModel) {
    116         // Before we start the bot, register our command handler
    117         // eslint-disable-next-line @typescript-eslint/no-misused-promises
    118         client.on("room.message", this.handleMessage.bind(this));
    119         // eslint-disable-next-line @typescript-eslint/no-misused-promises
    120         client.on("room.event", this.handleEvents.bind(this));
    121 
    122         // Now that everything is set up, start the bot. This will start the sync loop and run until killed.
    123         client.start().then(async () => {
    124             LogService.info("index", "Bot started!");
    125             // Send notice that bot is starting into both rooms
    126             await client.sendNotice(config.adminRoom, startUpMessage);
    127             await client.sendNotice(config.warningsRoom, startUpMessage);
    128         }).catch(console.error);
    129     }
    130 
    131     // eslint-disable-next-line @typescript-eslint/no-explicit-any
    132     private async handleEvents(roomId: string, ev: any): Promise<void> {
    133         // For now only handle reactions
    134         const event = new ReactionEvent(ev);
    135         if (event.isRedacted) return; // Ignore redacted events
    136         if (event.sender === await this.client.getUserId()) return; // Ignore ourselves
    137 
    138 
    139         try {
    140             await getReactionHandler(roomId, event, this.client, this.config, this.policyRoomHandler).handleReaction();
    141         } catch (e) {
    142             return;
    143             // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
    144             //LogService.error("index", `Error handling reaction: ${e}`);
    145         }
    146     }
    147 
    148     // This is the command handler we registered a few lines up
    149     // eslint-disable-next-line @typescript-eslint/no-explicit-any
    150     private async handleMessage(roomId: string, ev: any): Promise<void> {
    151         const event = new MessageEvent(ev);
    152         if (event.isRedacted) return; // Ignore redacted events that come through
    153         if (event.sender === await this.client.getUserId()) return; // Ignore ourselves
    154         if (event.messageType !== "m.text") return; // Ignore non-text messages
    155 
    156         if (roomId !== this.config.adminRoom && roomId !== this.config.warningsRoom) {
    157             await this.checkSpam(event, event.content?.body.trim() ?? "", roomId);
    158         }
    159     }
    160 
    161     // eslint-disable-next-line @typescript-eslint/no-explicit-any
    162     private async checkSpam(event: MessageEvent<MessageEventContent>, body: string, roomId: string) {
    163         // Check if spam
    164         const data = tf.tensor([body])
    165         const prediction: Tensor<Rank.R2> = this.model.predict(data) as Tensor<Rank.R2>;
    166         const prediction_data: number[][] = await prediction.array();
    167         LogService.info("index", `Checking: "${body}"`);
    168         LogService.info("index", `Prediction: ${prediction_data.toString()}`);
    169 
    170         const prediction_value = ((prediction_data[0] ?? [])[0] ?? 0);
    171         if (prediction_value > THRESHOLD) {
    172             const mxid = event.sender;
    173             const displayname = (await (this.client.getUserProfile(mxid) as Promise<MatrixProfileInfo>)).displayname ?? mxid;
    174             let alias = roomId;
    175             let roomname = roomId;
    176             let room_url = `matrix:roomid/${roomId.replace("!", "")}?via=${this.config.homeserver.replace("https://", "")}`;
    177             try {
    178                 alias = (await (this.client.getRoomStateEvent(roomId, "m.room.canonical_alias", "") as Promise<CanonicalAliasEventContent>)).alias;
    179                 roomname = alias;
    180                 room_url = `matrix:r/${alias.replace("#", "").replace("!", "")}`;
    181             } catch (e) {
    182                 LogService.debug("index", `Failed to get alias for ${roomId}`);
    183             }
    184             try {
    185                 roomname = (await (this.client.getRoomStateEvent(roomId, "m.room.name", "") as Promise<RoomNameEventContent>)).name;
    186             } catch (e) {
    187                 LogService.debug("index", `Failed to get name for ${roomId}`);
    188             }
    189 
    190             const html = `<blockquote>\n<p>${body}</p>\n</blockquote>\n<p>Above message was detected as spam. See json file for full event and use reactions to take action or no action.</p><p>It was sent by <a href="matrix:u/${mxid.replace("@", "")}">${displayname}</a> in <a href="${room_url}">${roomname}</a></p><p>Spam Score is: ${prediction_value.toFixed(3)}</p>\n`;
    191             const roominfo: { roomId: string; name: string | undefined; alias: string | undefined; } = {
    192                 roomId: roomId,
    193                 name: undefined,
    194                 alias: undefined
    195             };
    196             if (roomname !== roomId) {
    197                 roominfo["name"] = roomname;
    198             }
    199             if (alias !== roomId) {
    200                 roominfo["alias"] = alias;
    201             }
    202             const alert_event_id = await this.client.sendMessage(this.config.warningsRoom, {
    203                 body: htmlToText(html, { wordwrap: false }),
    204                 msgtype: "m.text",
    205                 format: "org.matrix.custom.html",
    206                 formatted_body: html,
    207                 "space.midnightthoughts.spam_score": prediction_value.toFixed(3),
    208                 "space.midnightthoughts.sending_user": mxid,
    209                 "space.midnightthoughts.sending_room": roominfo,
    210                 "space.midnightthoughts.event_id": event.eventId,
    211             });
    212 
    213             await this.client.unstableApis.addReactionToEvent(this.config.warningsRoom, alert_event_id, BAN_REACTION);
    214             await this.client.unstableApis.addReactionToEvent(this.config.warningsRoom, alert_event_id, KICK_REACTION);
    215             await this.client.unstableApis.addReactionToEvent(this.config.warningsRoom, alert_event_id, FALSE_POSITIVE_REACTION);
    216             const eventContent = Buffer.from(JSON.stringify(event.raw), 'utf8');
    217             const media = await this.client.uploadContent(eventContent, "application/json", "event.json");
    218             await this.client.sendMessage(this.config.warningsRoom, {
    219                 msgtype: "m.file",
    220                 body: "event.json",
    221                 filename: "event.json",
    222                 info: {
    223                     mimetype: "application/json",
    224                     size: eventContent.length,
    225                 },
    226                 url: media,
    227             })
    228         } else {
    229             // const textEvent = new MessageEvent<TextualMessageEventContent>(event.raw);
    230             //await this.client.unstableApis.addReactionToEvent(roomId, textEvent.eventId, "Classified Not Spam")
    231         }
    232     }
    233 }
    234 
    235 await Bot.createBot();