invitron/lib/InviteTracker.ts

A powerful Discord.js invite tracker with persistent storage, member analytics, vanity URL support, and comprehensive join monitoring system for Discord bots.

import {
  Client,
  Guild,
  GuildMember,
  Invite,
  User,
  PartialUser,
  GatewayIntentBits,
  PartialGuildMember,
} from "discord.js";
import { TrackerOptions } from "./index";
import { JSONStorage } from "./Storage";

export interface InviteJoinData {
  guildId: string;
  code: string;
  url: string | null;
  uses: number | null;
  maxUses: number | null;
  maxAge: number | null;
  createdTimestamp: number | null;
  inviter: User | PartialUser | null;
  totalInvites: number | null;
  type: "normal" | "vanity" | "audit" | "unknown";
}

interface InviteCodeData {
  uses: number;
  members: string[];
}

interface InviterData {
  total: number;
  codes: Record<string, InviteCodeData>;
}

interface PersistedData {
  [guildId: string]: {
    [inviterId: string]: InviterData;
  };
}

export class InviteTracker {
  private client: Client;
  private cache: Record<string, Record<string, number>> = {};
  private counts: PersistedData = {};
  private options: TrackerOptions;
  private storage?: JSONStorage;

  constructor(client: Client, options: TrackerOptions = {}) {
    this.client = client;
    this.options = options;

    // Initialize storage if enabled
    if (options.storage?.enabled) {
      this.storage = new JSONStorage(options.storage.path);
      this.counts = (this.storage.load() as PersistedData) || {};
    }

    // Ensure required intents are enabled
    if (
      !client.options.intents?.has(GatewayIntentBits.GuildInvites) ||
      !client.options.intents?.has(GatewayIntentBits.GuildMembers)
    ) {
      throw new Error(
        "InvitesTracker Error: Missing intents (GuildInvites, GuildMembers)!"
      );
    }

    // Preload invites from all guilds on client ready
    if (options.fetchGuilds) {
      this.client.once("ready", async () => {
        for (const [, guild] of this.client.guilds.cache) {
          await this.loadGuildInvites(guild);
        }
      });
    }

    // Track invite creation events
    this.client.on("inviteCreate", (invite) => {
      this.handleInviteCreate(invite);
      if (options.debug) console.log(`[DEBUG] Invite created: ${invite.code}`);
    });

    // Track invite deletion events
    this.client.on("inviteDelete", (invite) => {
      this.handleInviteDelete(invite);
      if (options.debug) console.log(`[DEBUG] Invite deleted: ${invite.code}`);
    });

    // Track member leave events
    this.client.on("guildMemberRemove", (member) => {
      this.handleMemberLeave(member);
    });
  }

  /** Persist current invite counts to storage */
  private persist() {
    if (this.storage) {
      this.storage.save(this.counts);
    }
  }

  /** Load all current invites for a guild into cache */
  private async loadGuildInvites(guild: Guild) {
    try {
      const invites = await guild.invites.fetch();
      this.cache[guild.id] = {};
      invites.forEach((inv) => {
        this.cache[guild.id][inv.code] = inv.uses ?? 0;
      });
    } catch {
      this.cache[guild.id] = {};
    }
  }

  /** Handle new invite creation and update cache */
  private handleInviteCreate(invite: Invite) {
    if (!invite.guild) return;
    this.cache[invite.guild.id] = this.cache[invite.guild.id] ?? {};
    this.cache[invite.guild.id][invite.code] = invite.uses ?? 0;
  }

  /** Handle invite deletion and update cache */
  private handleInviteDelete(invite: Invite) {
    if (!invite.guild) return;
    if (this.cache[invite.guild.id]) delete this.cache[invite.guild.id][invite.code];
  }

  /**
   * Handle member joining a guild
   * Updates inviter stats based on which invite was used
   */
  async handleMemberJoin(
    member: GuildMember,
    options: TrackerOptions
  ): Promise<InviteJoinData | null> {
    if (this.options.ignoreBots && member.user.bot) {
      if (this.options.debug) {
        console.log(`[DEBUG] Ignored bot join: ${member.user.tag}`);
      }
      return null;
    }

    try {
      const newInvites = await member.guild.invites.fetch();
      const oldInvites = this.cache[member.guild.id] ?? {};
      let usedInvite: Invite | null = null;

      // Determine which invite was used
      for (const [code, inv] of newInvites) {
        const before = oldInvites[code] ?? 0;
        const after = inv.uses ?? 0;
        if (after > before) {
          usedInvite = inv;
          break;
        }
      }

      // Update cache with current invite uses
      this.cache[member.guild.id] = { ...this.cache[member.guild.id] };
      newInvites.forEach((inv) => {
        this.cache[member.guild.id][inv.code] = inv.uses ?? 0;
      });

      // Process normal invite
      if (usedInvite) {
        const inviter = usedInvite.inviter ?? null;
        if (inviter) {
          const guildId = member.guild.id;
          const inviterId = inviter.id;
          const code = usedInvite.code;

          if (!this.counts[guildId]) this.counts[guildId] = {};
          if (!this.counts[guildId][inviterId]) {
            this.counts[guildId][inviterId] = { total: 0, codes: {} };
          }
          if (!this.counts[guildId][inviterId].codes[code]) {
            this.counts[guildId][inviterId].codes[code] = { uses: 0, members: [] };
          }

          this.counts[guildId][inviterId].total++;
          this.counts[guildId][inviterId].codes[code].uses++;
          this.counts[guildId][inviterId].codes[code].members.push(member.id);

          this.persist();
        }

        return {
          guildId: member.guild.id,
          code: usedInvite.code,
          url: usedInvite.url,
          uses: usedInvite.uses ?? null,
          maxUses: usedInvite.maxUses ?? null,
          maxAge: usedInvite.maxAge ?? null,
          createdTimestamp: usedInvite.createdTimestamp ?? null,
          inviter,
          totalInvites: inviter
            ? this.getUserInvites(member.guild.id, inviter.id)
            : null,
          type: "normal",
        };
      }

      // Handle vanity URL
      if (options.fetchVanity && member.guild.vanityURLCode) {
        return {
          guildId: member.guild.id,
          code: member.guild.vanityURLCode,
          url: `https://discord.gg/${member.guild.vanityURLCode}`,
          uses: null,
          maxUses: null,
          maxAge: null,
          createdTimestamp: null,
          inviter: null,
          totalInvites: null,
          type: "vanity",
        };
      }

      // Handle audit logs as fallback
      if (options.fetchAuditLogs) {
        try {
          const logs = await member.guild.fetchAuditLogs({ limit: 1, type: 28 });
          const entry = logs.entries.first();
          if (entry && entry.target?.id === member.id) {
            return {
              guildId: member.guild.id,
              code: "audit-log",
              url: null,
              uses: null,
              maxUses: null,
              maxAge: null,
              createdTimestamp: entry.createdTimestamp,
              inviter: entry.executor ?? null,
              totalInvites: entry.executor
                ? this.getUserInvites(member.guild.id, entry.executor.id)
                : null,
              type: "audit",
            };
          }
        } catch {
          // Ignore errors
        }
      }

      // Unknown join type
      return {
        guildId: member.guild.id,
        code: "unknown",
        url: null,
        uses: null,
        maxUses: null,
        maxAge: null,
        createdTimestamp: null,
        inviter: null,
        totalInvites: null,
        type: "unknown",
      };
    } catch {
      return null;
    }
  }

  /**
   * Handle member leaving a guild
   * Deducts invites if configured to do so
   */
  private handleMemberLeave(member: GuildMember | PartialGuildMember) {
    if (this.options.ignoreBots && member.user?.bot) return;

    if (this.options.deductOnLeave) {
      const guildId = member.guild.id;
      const guildData = this.counts[guildId];
      if (!guildData) return;

      for (const [inviterId, inviterData] of Object.entries(guildData)) {
        for (const [code, codeData] of Object.entries(inviterData.codes)) {
          const index = codeData.members.indexOf(member.id);
          if (index !== -1) {
            // Remove member from invite record
            codeData.members.splice(index, 1);

            // Deduct usage count
            if (codeData.uses > 0) codeData.uses--;

            // Deduct from total invites
            if (inviterData.total > 0) inviterData.total--;

            if (this.options.debug) {
              console.log(
                `[DEBUG] Member ${member.id} left, deducted invite from inviter ${inviterId}, code ${code}`
              );
            }

            // Save and exit after deduction
            this.persist();
            return;
          }
        }
      }
    }

    this.persist();
  }

  /** Clean invalid members from guild invite counts */
  async cleanGuild(guildId: string) {
    const guild = this.client.guilds.cache.get(guildId);
    if (!guild) return;

    const members = await guild.members.fetch();
    const validIds = new Set(members.map((m) => m.id));

    if (this.counts[guildId]) {
      for (const [inviterId, inviterData] of Object.entries(this.counts[guildId])) {
        for (const [code, codeData] of Object.entries(inviterData.codes)) {
          codeData.members = codeData.members.filter((m) => validIds.has(m));
          codeData.uses = codeData.members.length;
        }
        inviterData.total = Object.values(inviterData.codes).reduce((acc, c) => acc + c.uses, 0);
      }
    }

    this.persist();
  }

  /** Get total invites for a user in a guild */
  getUserInvites(guildId: string, userId: string): number {
    return this.counts[guildId]?.[userId]?.total ?? 0;
  }

  /** Reset all invite data for a guild */
  resetGuild(guildId: string) {
    this.counts[guildId] = {};
    this.persist();
  }

  /** Get the top inviter in a guild */
  getTopInviter(guildId: string): { userId: string; invites: number } | null {
    const guildInvites = this.counts[guildId];
    if (!guildInvites) return null;
    const [userId, data] =
      Object.entries(guildInvites).sort((a, b) => b[1].total - a[1].total)[0] || [];
    return userId ? { userId, invites: data.total } : null;
  }

  /** Get the leaderboard for a guild, optionally limited */
  getLeaderboard(
    guildId: string,
    limit = 10
  ): { userId: string; invites: number }[] {
    const guildInvites = this.counts[guildId];
    if (!guildInvites) return [];
    return Object.entries(guildInvites)
      .sort((a, b) => b[1].total - a[1].total)
      .slice(0, limit)
      .map(([userId, data]) => ({ userId, invites: data.total }));
  }
}