diff --git a/README.md b/README.md index f4b7101..785c155 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,16 @@ This metric can be enabled with command line flag `--words-counter` `--keywords-file`, `-k` - path to yaml file with keywords and patterns (see [keywords.yml.example](./keywords.yml.example)) +`--watch-file`, `-w` - watch for keywords file updates and reload keywords configuration in runtime + +`--include-peers`, `-i` - comma-separated list of `peer.id`s to gather metrics from. +if set, only specified peers will be exposed in metrics. +can be specified multiple times. can not be used along with `--exclude-peers` + +`--exclude-peers` `x` - comma-separated list of `peer.id`s to exclude from metrics. +if set, specified peers will not be exposed in metrics. +can be specified multiple times. can not be used along with `--include-peers` + ## Environment Variables `API_ID` - Telegram api id used for mtproto connection (see [mtcute.dev](https://mtcute.dev/guide/intro/sign-in.html)) diff --git a/docker-compose.yaml b/docker-compose.yaml index e1ce63c..7e5edc8 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -7,4 +7,11 @@ services: env_file: - .env volumes: - - ./bot-data:/app/bot-data \ No newline at end of file + - ./bot-data:/app/bot-data + - ./keywords.yml:/app/keywords.yml + command: + - "--keywords-file" + - "/app/keywords.yml" + - "--watch-file" + ports: + - 127.0.0.1:9669:9669 diff --git a/package.json b/package.json index 6d50d83..fb93db9 100644 --- a/package.json +++ b/package.json @@ -1,7 +1,7 @@ { "name": "mtproto_exporter", "type": "module", - "version": "1.0.1", + "version": "1.1.0", "packageManager": "pnpm@10.6.5", "license": "MIT", "scripts": { diff --git a/src/config.ts b/src/config.ts new file mode 100644 index 0000000..f2c93e1 --- /dev/null +++ b/src/config.ts @@ -0,0 +1,78 @@ +import type { OptionDefinition } from "command-line-args"; +import type { KeywordLike } from "./keywords.js"; +import { readFile } from "node:fs/promises"; +import cmdline from "command-line-args"; +import yaml from "js-yaml"; + +export interface Configuration { + bindHost: string; + port: number; + wordsCounter: boolean; + keywordsFile: string; + watchFile: boolean; + includePeers?: number[]; + excludePeers?: number[]; + keywords?: KeywordLike[]; +} + +const optionDefinitions: OptionDefinition[] = [ + { name: "bind-host", alias: "b", type: String, defaultValue: "0.0.0.0" }, + { name: "port", alias: "p", type: Number, defaultValue: 9669 }, + { name: "words-counter", type: Boolean, defaultValue: false }, + { name: "keywords-file", alias: "k", type: String }, + { name: "watch-file", alias: "w", type: Boolean, defaultValue: false }, + { name: "include-peers", alias: "i", type: String, multiple: true }, + { name: "exclude-peers", alias: "x", type: String, multiple: true }, +]; + +const cli = cmdline(optionDefinitions); + +const config: Configuration = { + bindHost: cli["bind-host"], + port: cli.port, + wordsCounter: cli["words-counter"], + keywordsFile: cli["keywords-file"], + watchFile: cli["watch-file"], + keywords: cli["keywords-file"] ? await readKeywords(cli["keywords-file"]) : undefined, +}; + +if (cli["include-peers"] && cli["exclude-peers"]) { + throw new Error("Conflicting configuration: --include-peers and --exclude-peers can't be both specified at the same time."); +} + +if (cli["include-peers"]) { + config.includePeers = []; + for (const o of cli["include-peers"] as string[]) { + config.includePeers.push(...o.split(",").map(i => parseInt(i)).filter(i => !isNaN(i))); + } +} + +if (cli["exclude-peers"]) { + config.excludePeers = []; + for (const o of cli["exclude-peers"] as string[]) { + config.excludePeers.push(...o.split(",").map(i => parseInt(i)).filter(i => !isNaN(i))); + } +} + +export async function readKeywords(filePath: string): Promise { + const doc = yaml.load(await readFile(filePath, "utf8")) as { keywords?: any[] }; + + if (doc.keywords && doc.keywords.constructor.name === "Array") { + const keywords: KeywordLike[] = []; + for (const item of doc.keywords) { + if (typeof item === "string") { + keywords.push(item); + } else if (typeof item === "object" && item.name && item.pattern) { + keywords.push({ + name: item.name, + pattern: new RegExp(item.pattern, "gi"), + }); + } + } + return keywords; + } else { + throw new Error("Keywords file format error: no 'keywords' property, or not an array."); + } +} + +export { config }; diff --git a/src/filters.ts b/src/filters.ts new file mode 100644 index 0000000..cef2a82 --- /dev/null +++ b/src/filters.ts @@ -0,0 +1,21 @@ +import { filters } from "@mtcute/dispatcher"; +import { Configuration } from "./config.js"; + +export function peersConfigBoolFilter(conf: Configuration, peerId: number) { + if(conf.excludePeers && conf.excludePeers.indexOf(peerId) !== -1) { + return false; + } + if(conf.includePeers && conf.includePeers.indexOf(peerId) === -1) { + return false; + } + return true; +} + +export function peersConfigFilter(conf: Configuration) { + if(conf.excludePeers) { + return filters.not(filters.chatId(conf.excludePeers)); + } else if (conf.includePeers) { + return filters.chatId(conf.includePeers); + } + return filters.any; +} \ No newline at end of file diff --git a/src/keywords.ts b/src/keywords.ts index 3dd3d6c..3340cf6 100644 --- a/src/keywords.ts +++ b/src/keywords.ts @@ -1,6 +1,8 @@ import type { Dispatcher } from "@mtcute/dispatcher"; import { PropagationAction } from "@mtcute/dispatcher"; import { Counter } from "prom-client"; +import { peersConfigFilter } from "./filters.js"; +import { config } from "./config.js"; interface KeywordPattern { name: string; @@ -15,7 +17,7 @@ export function newWordsCounter(dp: Dispatcher) { help: "Number of words in messages since exporter startup", labelNames: ["peerId", "word"], }); - dp.onNewMessage(async (msg) => { + dp.onNewMessage(peersConfigFilter(config), async (msg) => { const words = msg.text.toLowerCase().split(" "); for (const w of words) { counter.inc({ @@ -28,34 +30,47 @@ export function newWordsCounter(dp: Dispatcher) { return counter; } -export function newKeywordsCounter(dp: Dispatcher, keywords: KeywordLike[]) { - const counter = new Counter({ - name: "messenger_dialog_keywords_count", - help: "Number of keywords found in messages since exporter startup", - labelNames: ["peerId", "keyword"], - }); - dp.onNewMessage(async (msg) => { - for (const kw of keywords) { - let count; - let kwname; - if (typeof kw === "string") { - const words = msg.text.toLowerCase().split(" "); - count = words.filter(w => w === kw).length; - kwname = kw; - } else { - count = (msg.text.match(kw.pattern) || []).length; - kwname = kw.name; - } - if (count === 0) { - continue; - } - counter.inc({ - peerId: msg.chat.id, - keyword: kwname, - }, count); - } - return PropagationAction.Continue; - }); +export class KeywordsCounter extends Counter { + private _dp: Dispatcher; + private _keywords: KeywordLike[]; + constructor(dp: Dispatcher, keywords: KeywordLike[] = []) { + super({ + name: "messenger_dialog_keywords_count", + help: "Number of keywords found in messages since exporter startup", + labelNames: ["peerId", "keyword"], + }); + this._dp = dp; + this._keywords = keywords; - return counter; + dp.onNewMessage(peersConfigFilter(config), async (msg) => { + for (const kw of this._keywords) { + let count; + let kwname; + if (typeof kw === "string") { + const words = msg.text.toLowerCase().split(" "); + count = words.filter(w => w === kw).length; + kwname = kw; + } else { + count = (msg.text.match(kw.pattern) || []).length; + kwname = kw.name; + } + if (count === 0) { + continue; + } + this.inc({ + peerId: msg.chat.id, + keyword: kwname, + }, count); + } + return PropagationAction.Continue; + }); + } + + public get keywords() { + return this._keywords; + } + + public setKeywords(keywords: KeywordLike[]) { + this._keywords = keywords; + } } diff --git a/src/main.ts b/src/main.ts index 447fdd2..7de6300 100644 --- a/src/main.ts +++ b/src/main.ts @@ -1,54 +1,19 @@ -import type { OptionDefinition } from "command-line-args"; -import type { KeywordLike } from "./keywords.js"; import fs from "node:fs"; import { Dispatcher } from "@mtcute/dispatcher"; import { TelegramClient } from "@mtcute/node"; -import cmdline from "command-line-args"; -import yaml from "js-yaml"; import { collectDefaultMetrics, Registry } from "prom-client"; +import { config, readKeywords } from "./config.js"; import * as env from "./env.js"; import * as metrics from "./metrics.js"; import MetricsServer from "./server.js"; -const optionDefinitions: OptionDefinition[] = [ - { name: "bind-host", alias: "b", type: String, defaultValue: "0.0.0.0" }, - { name: "port", alias: "p", type: Number, defaultValue: 9669 }, - { name: "words-counter", type: Boolean, defaultValue: false }, - { name: "keywords-file", alias: "k", type: String }, -]; - -const cli = cmdline(optionDefinitions); - -const keywords: KeywordLike[] = []; -if (cli["keywords-file"]) { - if (!fs.existsSync(cli["keywords-file"])) { - throw new Error("--keywords-file set, but file does not exist."); - } - const doc = yaml.load(fs.readFileSync(cli["keywords-file"], "utf8")) as { keywords?: any[] }; - - if (doc.keywords && doc.keywords.constructor.name === "Array") { - for (const item of doc.keywords) { - if (typeof item === "string") { - keywords.push(item); - } else if (typeof item === "object" && item.name && item.pattern) { - keywords.push({ - name: item.name, - pattern: new RegExp(item.pattern, "gi"), - }); - } - } - } else { - throw new Error("Keywords file format error: no 'keywords' property, or not an array."); - } -} - const registry = new Registry(); collectDefaultMetrics({ register: registry }); const server = new MetricsServer(registry); -server.listen(cli["bind-host"], cli.port); +server.listen(config.bindHost, config.port); const tg = new TelegramClient({ apiId: env.API_ID, @@ -70,10 +35,26 @@ registry.registerMetric(metrics.newStaticPeerInfoGauge(tg)); registry.registerMetric(metrics.newUnreadCountGauge(tg)); registry.registerMetric(metrics.newMessagesCounter(dp)); -if (keywords.length > 0) { - registry.registerMetric(metrics.newKeywordsCounter(dp, keywords)); +if (config.keywords) { + const counter = new metrics.KeywordsCounter(dp, config.keywords); + registry.registerMetric(counter); + + if (config.watchFile) { + fs.watchFile(config.keywordsFile, async (curr, prev) => { + if (curr.mtimeMs === prev.mtimeMs) { + return; + } + console.log("[watch-file] Keywords file was updated. Re-reading keywords configuration..."); + try { + config.keywords = await readKeywords(config.keywordsFile); + counter.setKeywords(config.keywords); + } catch (e) { + console.error("Failed to read keywords file", config.keywordsFile, e); + } + }); + } } -if (cli["words-counter"]) { +if (config.wordsCounter) { registry.registerMetric(metrics.newWordsCounter(dp)); } diff --git a/src/metrics.ts b/src/metrics.ts index ed2b376..527cb68 100644 --- a/src/metrics.ts +++ b/src/metrics.ts @@ -3,7 +3,9 @@ import type { TelegramClient } from "@mtcute/node"; import { PropagationAction } from "@mtcute/dispatcher"; import { Counter, Gauge } from "prom-client"; -import { newKeywordsCounter, newWordsCounter } from "./keywords.js"; +import { KeywordsCounter, newWordsCounter } from "./keywords.js"; +import { config } from "./config.js"; +import { peersConfigBoolFilter, peersConfigFilter } from "./filters.js"; function newMessagesCounter(dp: Dispatcher) { const counter = new Counter({ @@ -11,7 +13,8 @@ function newMessagesCounter(dp: Dispatcher) { help: "Messages count since exporter startup", labelNames: ["peerId"], }); - dp.onNewMessage(async (msg) => { + + dp.onNewMessage(peersConfigFilter(config), async (msg) => { counter.inc({ peerId: msg.chat.id, }); @@ -28,6 +31,10 @@ function newStaticPeerInfoGauge(tg: TelegramClient) { collect: async () => { gauge.reset(); for await (const d of tg.iterDialogs()) { + if(!peersConfigBoolFilter(config, d.peer.id)) { + continue; + } + gauge.set({ peerId: d.peer.id, peerType: d.peer.type, @@ -47,6 +54,10 @@ function newUnreadCountGauge(tg: TelegramClient) { collect: async () => { gauge.reset(); for await (const d of tg.iterDialogs()) { + if(!peersConfigBoolFilter(config, d.peer.id)) { + continue; + } + gauge.set({ peerId: d.peer.id, }, d.unreadCount); @@ -57,7 +68,7 @@ function newUnreadCountGauge(tg: TelegramClient) { } export { - newKeywordsCounter, + KeywordsCounter, newMessagesCounter, newStaticPeerInfoGauge, newUnreadCountGauge,