Skip to content

Commit

Permalink
Allow dumping attachments
Browse files Browse the repository at this point in the history
  • Loading branch information
haykam821 committed Sep 30, 2023
1 parent c43c7cb commit 9461764
Show file tree
Hide file tree
Showing 5 changed files with 151 additions and 35 deletions.
2 changes: 1 addition & 1 deletion package.json
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,7 @@
"devDependencies": {
"@types/debug": "^4.1.7",
"@types/fs-extra": "^9.0.13",
"@types/node": "^17.0.23",
"@types/node": "^20.8.0",
"@typescript-eslint/eslint-plugin": "^6.7.3",
"@typescript-eslint/parser": "^6.7.3",
"eslint": "^8.43.0",
Expand Down
65 changes: 31 additions & 34 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,27 +3,18 @@
import { CHILD_CHANNEL_EMOJI, GUILD_OWNER_EMOJI, INFO_HEADER_EMOJI, JOIN_MESSAGE_EMOJI, MESSAGE_EMOJI, MESSAGE_WITH_ATTACHMENT_EMOJI, NAME_CHANGE_MESSAGE_EMOJI, NO_PERMISSION_EMOJI, PIN_MESSAGE_EMOJI, REPLY_MESSAGE_EMOJI, THREAD_MESSAGE_EMOJI, TTS_MESSAGE_EMOJI, UNKNOWN_MESSAGE_EMOJI } from "./emoji";
import djs, { BaseChannel, BaseGuildTextChannel, CategoryChannel, Client, ClientOptions, Collection, DMChannel, DiscordAPIError, Guild, GuildChannel, Message, MessageReaction, MessageType, NewsChannel, Snowflake, TextChannel, ThreadChannel } from "discord.js";

import { AttachmentStore } from "./utils/attachment-store";
import { GuildPaths } from "./utils/guild-paths";
import { WriteStream } from "node:fs";
import cli from "caporal";
import debug from "debug";
import fs from "fs-extra";
import { getChannelType } from "./utils/channel-type";
import { log } from "./utils/log";
import open from "open";
import path from "node:path";
import readPkg from "read-pkg";

/**
* The timestamp used in part of the dump's path.
*/
const dumpDate = Date.now().toString();

// Set up logging with debug module
const log = {
dumper: debug("discord-dumper:dumper"),
path: debug("discord-dumper:path"),
prepare: debug("discord-dumper:prepare"),
};

/**
* Gets a string representing a vessel name.
* @param vessel The vessel to get the display name of.
Expand Down Expand Up @@ -73,7 +64,9 @@ function emojiName(reaction: MessageReaction): string {
async function dumpHierarchy(guild: Guild): Promise<void> {
if (!(guild instanceof Guild)) return;

const hierarchyPath = path.resolve(`./dumps/${guild.id}/${dumpDate}/member_hierarchy.txt`);
const paths = new GuildPaths(null);
const hierarchyPath = paths.getMemberHierarchy();

await fs.ensureFile(hierarchyPath);
const hierarchyStream = fs.createWriteStream(hierarchyPath);

Expand Down Expand Up @@ -107,9 +100,10 @@ async function dumpHierarchy(guild: Guild): Promise<void> {
/**
* Writes a single message's information to a stream.
* @param dumpStream The stream to write the message to.
* @param attachments The attachment store to add attachments to.
* @param message The message itself.
*/
function dumpMessage(dumpStream: WriteStream, message: Message): void {
function dumpMessage(dumpStream: WriteStream, attachments: AttachmentStore, message: Message): void {
const dumpMessage_ = [
` ${message.id.padStart(18)} `,
`[${message.createdAt.toLocaleString()}] `,
Expand Down Expand Up @@ -173,6 +167,8 @@ function dumpMessage(dumpStream: WriteStream, message: Message): void {
}
dumpMessage_.push(` ${message.attachments.map(attachment => attachment.url)
.join(" ")}`);

attachments.addAll(message.attachments);
} else {
let emoji;

Expand All @@ -198,28 +194,26 @@ function dumpMessage(dumpStream: WriteStream, message: Message): void {
dumpStream.write(dumpMessage_.join("") + "\n");
}

/**
* Ensures and gets the path to where dumps should be stored for a channel at the given dump time.
* @param channel The channel to get the relevant dump path for.
* @returns - The path where the dumps should be stored for a channel at the given dump time.
*/
async function channelify(channel: BaseChannel): Promise<string> {
const guildName = channel instanceof GuildChannel || channel instanceof ThreadChannel ? channel.guild.id : "guildless";
const pathToChannel = path.resolve(`./dumps/${guildName}/${dumpDate}/channels/${channel.id}.txt`);

await fs.ensureFile(pathToChannel);

return pathToChannel;
}

/**
* Dumps a channel with its basic information and its messages.
* @param channel The channel to dump.
* @param [shouldDumpMessages=true] Whether to dump messages or not.
* @param [shouldDumpAttachments=true] Whether to dump attachments or not.
*/
async function dump(channel: BaseChannel, shouldDumpMessages = true) {
const dumpPath = await channelify(channel);
const dumpStream = fs.createWriteStream(dumpPath);
async function dump(channel: BaseChannel, shouldDumpMessages = true, shouldDumpAttachments = true) {
const paths = new GuildPaths(channel);

const channelPath = paths.getChannel();
const attachmentsPath = paths.getAttachments();

await fs.ensureFile(channelPath);
const dumpStream = fs.createWriteStream(channelPath);

if (shouldDumpAttachments) {
await fs.ensureDir(attachmentsPath);
}

const attachments = new AttachmentStore(attachmentsPath);

dumpStream.write([
INFO_HEADER_EMOJI + ` Name: ${displayName(channel)} (${getChannelType(channel)})`,
Expand Down Expand Up @@ -270,7 +264,7 @@ async function dump(channel: BaseChannel, shouldDumpMessages = true) {
} else {
oldestDumped = fetches.last()?.id;
for (const messageToDump of fetches.values()) {
await dumpMessage(dumpStream, messageToDump);
await dumpMessage(dumpStream, attachments, messageToDump);
}
}
} catch (error) {
Expand All @@ -285,6 +279,8 @@ async function dump(channel: BaseChannel, shouldDumpMessages = true) {
}
}

await attachments.dump(shouldDumpAttachments);

return dumpStream.end();
}

Expand Down Expand Up @@ -437,10 +433,10 @@ async function likeActuallyDump(vessel: Vessel, argv: Record<string, unknown>) {
await dumpHierarchy(vessel);
}
for (const channel of vessel.channels.cache) {
await dump(channel[1], argv.dumpMessages as boolean);
await dump(channel[1], argv.dumpMessages as boolean, argv.dumpAttachments as boolean);
}
} else if (vessel instanceof djs.BaseChannel) {
await dump(vessel, argv.dumpMessages as boolean);
await dump(vessel, argv.dumpMessages as boolean, argv.dumpAttachments as boolean);
}
}

Expand All @@ -451,6 +447,7 @@ cli
.option("--bypass [bypass]", "Uses the bypass, if it exists.", cli.BOOLEAN, true)
.option("--hierarchy [hierarchy]", "Dumps the role/member hierarchy of a guild.", cli.BOOLEAN, true)
.option("--dumpMessages [dumpMessages]", "Dumps the message history of channels.", cli.BOOLEAN, true)
.option("--dumpAttachments [dumpAttachments]", "Dumps attachments to local files.", cli.BOOLEAN, true)
.option("--path <path>", "The directory to store dumps in.", cli.STRING, "./dumps")
.argument("<id>", "The ID of the guild/channel/category/DM channel to dump.")
.argument("[context]", "The context of the ID.", Object.keys(contexts), "infer")
Expand Down
65 changes: 65 additions & 0 deletions src/utils/attachment-store.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
import { Attachment, Collection, Snowflake } from "discord.js";

import { Writable } from "node:stream";
import { createWriteStream } from "fs-extra";
import { log } from "./log";
import { resolve as resolvePath } from "node:path";

const EXTENSION_PATTERN = /([^./]+)$/;

export class AttachmentStore {
private readonly attachments = new Map<Snowflake, string>();
private readonly basePath: string;

constructor(basePath: string) {
this.basePath = basePath;
}

private getExtension(url: URL): string {
const match = url.pathname.match(EXTENSION_PATTERN);

if (match === null) {
log.attachments("Failed to find extension for attachment '%s'.", url);
return "";
}

return "." + match[1];
}

async dump(shouldDumpAttachments: boolean): Promise<void> {
if (!shouldDumpAttachments) {
log.attachments("Not dumping %d attachments.", this.attachments.size);
return;
}

log.attachments("Dumping %d attachments.", this.attachments.size);

for (const [id, urlString] of this.attachments) {
const url = new URL(urlString);

const extension = this.getExtension(url);
const path = resolvePath(this.basePath, `${id}${extension}`);

const response = await fetch(url);

if (response.body) {
const fileStream = createWriteStream(path);
response.body.pipeTo(Writable.toWeb(fileStream));
}

log.attachments("Dumped attachment '%s' to '%s'.", urlString, path);
}

return;
}

private add(id: Snowflake, attachment: Attachment): void {
this.attachments.set(id, attachment.proxyURL);
}

addAll(attachments: Collection<Snowflake, Attachment>): void {
for (const [id, attachment] of attachments) {
this.add(id, attachment);
}
}
}
46 changes: 46 additions & 0 deletions src/utils/guild-paths.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,46 @@
import { BaseChannel, GuildChannel, ThreadChannel } from "discord.js";

import { resolve } from "node:path";

/**
* The timestamp used in part of the dump's path.
*/
const dumpDate = Date.now().toString();

function getGuildIdPart(channel: BaseChannel | null): string {
if (channel instanceof GuildChannel || channel instanceof ThreadChannel) {
return channel.guild.id;
}

return "guildless";
}

export class GuildPaths {
private readonly channel: BaseChannel | null;
private readonly base: string;

constructor(channel: BaseChannel | null) {
this.channel = channel;
this.base = resolve(`./dumps/${getGuildIdPart(channel)}/${dumpDate}`);
}

private resolve(...segments: string[]) {
return resolve(this.base, ...segments);
}

getAttachments(): string {
return this.resolve("./attachments");
}

getChannel(): string {
if (this.channel === null) {
throw new Error("Cannot get channel path in non-channel context");
}

return this.resolve(`./channels/${this.channel.id}.txt`);
}

getMemberHierarchy(): string {
return this.resolve("./member_hierarchy.txt");
}
}
8 changes: 8 additions & 0 deletions src/utils/log.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
import debug from "debug";

export const log = {
attachments: debug("discord-dumper:attachments"),
dumper: debug("discord-dumper:dumper"),
path: debug("discord-dumper:path"),
prepare: debug("discord-dumper:prepare"),
};

0 comments on commit 9461764

Please sign in to comment.