auron/src/api/atproto.ts

Interact with your ATProto labeler from your terminal

import {
  AppBskyActorDefs,
  AtpAgent,
  ToolsOzoneModerationDefs,
  ToolsOzoneModerationQueryEvents,
  ToolsOzoneModerationQueryStatuses,
  ComAtprotoServerCreateSession,
} from "@atproto/api";
import { chunkArray } from "@atproto/common";
import { EventEmitter } from "events";
import * as readline from "readline";
import { getGlobalSpinner } from "../utils/loader";

let agent: AtpAgent;

const prompt2FA = (): Promise<string> => {
  return new Promise((resolve) => {
    const spinner = getGlobalSpinner();
    
    // Stop the spinner to allow clean input
    if (spinner) {
      spinner.stop();
    }
    
    const rl = readline.createInterface({
      input: process.stdin,
      output: process.stdout,
    });
    
    rl.question("Enter your 2FA code: ", (answer) => {
      rl.close();
      
      // Restart the spinner
      if (spinner) {
        spinner.start();
      }
      
      resolve(answer.trim());
    });
  });
};

export const getAgent = async () => {
  if (agent) return agent;
  if (!process.env.SERVICE_URL) {
    throw new Error("SERVICE_URL env var is required");
  }
  if (!process.env.SERVICE_DID) {
    throw new Error("SERVICE_DID env var is required");
  }

  agent = new AtpAgent({ service: process.env.SERVICE_URL });
  // @ts-ignore
  agent.configureProxy(process.env.SERVICE_DID!);
  
  try {
    await agent.login({
      identifier: process.env.USERNAME!,
      password: process.env.PASSWORD!,
    });
  } catch (error) {
    if (error instanceof ComAtprotoServerCreateSession.AuthFactorTokenRequiredError) {
      const authFactorToken = await prompt2FA();
      await agent.login({
        identifier: process.env.USERNAME!,
        password: process.env.PASSWORD!,
        authFactorToken,
      });
    } else {
      throw error;
    }
  }
  
  return agent;
};

export const getQueueItems = async (
  {
    reviewState,
    maxCount = 500,
    takendown,
    cursor,
  }: {
    reviewState?: string;
    maxCount?: number;
    cursor?: string;
    takendown?: boolean;
  },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  let nextCursor = cursor;
  let counter = 0;
  const result: ToolsOzoneModerationQueryStatuses.OutputSchema = {
    cursor,
    subjectStatuses: [],
  };

  do {
    try {
      const params: ToolsOzoneModerationQueryStatuses.QueryParams = {
        cursor: nextCursor,
        limit: Math.min(100, maxCount),
        reviewState: reviewState || ToolsOzoneModerationDefs.REVIEWOPEN,
      };
      if (takendown !== undefined) {
        params.takendown = takendown;
      }
      const { data } = await agent.tools.ozone.moderation.queryStatuses(params);
      nextCursor = data.cursor;
      result.subjectStatuses.push(...data.subjectStatuses);
      emitter.emit("update", {
        maxCount,
        nextCursor,
        subjectCount: result.subjectStatuses.length,
      });
    } catch (err) {
      console.error(err);
      break;
    }

    // Every 5th request, wait for 500ms to avoid potential rate limiting
    if (counter % 5) {
      await new Promise((resolve) => setTimeout(resolve, 500));
    }

    counter++;
  } while (
    nextCursor &&
    (!maxCount || result.subjectStatuses.length < maxCount)
  );

  return { ...result, cursor: nextCursor };
};

export const getEvents = async (
  {
    types,
    createdAfter,
    createdBefore,
    cursor,
  }: {
    types?: string[];
    createdAfter?: string;
    createdBefore?: string;
    cursor?: string;
  },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  let nextCursor = cursor;
  let counter = 0;
  const result: ToolsOzoneModerationQueryEvents.OutputSchema = {
    cursor,
    events: [],
  };

  do {
    try {
      const { data } = await agent.tools.ozone.moderation.queryEvents({
        types,
        limit: 100,
        createdAfter,
        createdBefore,
        cursor: nextCursor,
      });
      nextCursor = data.cursor;
      result.events.push(...data.events);
      emitter.emit("update", {
        nextCursor,
        eventCount: result.events.length,
      });
    } catch (err) {
      console.error(err);
      break;
    }

    // Every 5th request, wait for 500ms to avoid potential rate limiting
    if (counter % 5) {
      await new Promise((resolve) => setTimeout(resolve, 500));
    }

    counter++;
  } while (nextCursor);

  return { ...result, cursor: nextCursor };
};

export const getRepos = async (
  { dids }: { dids: string[] },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  const repos: ToolsOzoneModerationDefs.RepoViewDetail[] = [];

  for (const chunk of chunkArray(dids, 100)) {
    try {
      const { data } = await agent.tools.ozone.moderation.getRepos({
        dids: chunk,
      });
      repos.push(
        ...data.repos.filter((r) =>
          ToolsOzoneModerationDefs.isRepoViewDetail(r)
        )
      );

      emitter.emit("update", {
        total: dids.length,
        repoCount: repos.length,
      });
    } catch (err) {
      console.log(`Error fetching repos for ${chunk}`);
    }
  }

  return repos;
};

export const getProfiles = async (
  { dids }: { dids: string[] },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  const profiles: AppBskyActorDefs.ProfileViewDetailed[] = [];

  let i = 0;
  for (const chunk of chunkArray(dids, 25)) {
    if (i % 10 === 0) {
      await new Promise((res) => setTimeout(res, 1000));
    }
    try {
      const { data } = await agent.app.bsky.actor.getProfiles({
        actors: chunk,
      });
      profiles.push(...data.profiles);

      emitter.emit("update", {
        total: dids.length,
        profileCount: profiles.length,
      });
    } catch (err) {
      console.log(`Error fetching profiles for ${chunk}`);
    }
  }

  return profiles;
};

export const getRecords = async (
  { uris }: { uris: string[] },
  emitter: EventEmitter
) => {
  const agent = await getAgent();
  const records: ToolsOzoneModerationDefs.RecordViewDetail[] = [];

  for (const chunk of chunkArray(uris, 100)) {
    try {
      const { data } = await agent.tools.ozone.moderation.getRecords({
        uris: chunk,
      });
      records.push(
        ...data.records.filter((r) =>
          ToolsOzoneModerationDefs.isRecordViewDetail(r)
        )
      );

      emitter.emit("update", {
        total: uris.length,
        recordCount: records.length,
      });
    } catch (err) {
      console.log(`Error fetching records for ${chunk}`);
    }
  }

  return records;
};