neo-scraper/src/engines/Mastodon.ts

import {
  ScrapeEngineBase,
  ScrapeResult,
  ScrapedPost,
  ScrapeEngineFeature,
  ContentType,
} from "../ScrapeEngine.js";

interface Attachment {
  url: string;          // original URL, ?sensitive=true stripped
  previewUrl?: string;  // local cached preview (for video posters)
  contentType: ContentType;
  sensitive: boolean;
}

/**
 * Handles Mastodon-compatible instances, including forks such as baraag.net.
 *
 * On a Mastodon post page, media is rendered as:
 *   <div class="media-gallery__item">
 *     <a class="media-gallery__item-thumbnail" href=".../original/....png">
 *       <img src=".../small/....png">
 *     </a>
 *   </div>
 *
 * The anchor href is the original (often federated) file. The <img> is a
 * scaled-down local cache. We read the anchor href, never the img src.
 *
 * Multi-image posts repeat the same structure; we return one ScrapedPost
 * per anchor, matching how Pixiv/Twitter engines behave.
 */
export default class Mastodon extends ScrapeEngineBase {
  name = "mastodon";
  features: ScrapeEngineFeature[] = ["content", "rating", "source"];
  notes = [
    "Uses the original attachment URL from the media gallery anchor,",
    "not the scaled-down preview image.",
    "Defaults rating to 'safe'; bumps to 'explicit' when the",
    "attachment is marked sensitive in the source.",
  ];

  // Explicit hosts we claim. Other Mastodon instances are picked up by the
  // URL-shape check in canImport().
  supportedHosts = ["baraag.net"];

  canImport(url: Location): boolean {
    if (this.supportedHosts.indexOf(url.host.replace("www.", "")) !== -1) {
      return true;
    }
    // Generic Mastodon status URL shape: /@user/<numeric id>
    return /^\/@[^/]+\/\d+\/?$/.test(url.pathname);
  }

  scrapeDocument(document: Document): ScrapeResult {
    const result = new ScrapeResult(this.name);

    const attachments = this.findAttachments(document);
    this.log(`Found ${attachments.length} attachment(s).`);

    for (const att of attachments) {
      const post = new ScrapedPost();
      post.pageUrl = document.location.href;
      post.contentUrl = att.url;
      post.contentType = att.contentType;
      post.sources = [document.location.href];
      post.referrer = document.location.origin;

      // Baraag is an adult-focused instance: default to safe,
      // and bump to explicit when the source flags the attachment.
      post.rating = att.sensitive ? "explicit" : "safe";

      // Cached preview as a poster frame for videos.
      if (att.contentType === "video" && att.previewUrl) {
        post.extraContentUrl = att.previewUrl;
      }

      result.tryAddPost(post);
    }

    if (attachments.length === 0) {
      this.log("No media attachments found on this page.");
    }

    return result;
  }

  private findAttachments(document: Document): Attachment[] {
    const out: Attachment[] = [];

    // Primary selector, confirmed against baraag's rendered HTML for both
    // single- and multi-image posts.
    let anchors = Array.from(
      document.querySelectorAll<HTMLAnchorElement>("a.media-gallery__item-thumbnail[href]")
    );

    // Looser fallback for themes that drop the class.
    if (anchors.length === 0) {
      anchors = Array.from(
        document.querySelectorAll<HTMLAnchorElement>("a[href*='/media_attachments/']")
      );
    }

    for (const a of anchors) {
      const rawHref = a.getAttribute("href");
      if (!rawHref) continue;

      const img = a.querySelector("img");
      const preview = img?.currentSrc || img?.src;

      const { url, sensitive } = this.normalizeUrl(rawHref);
      const finalUrl = url || this.largestFromSrcset(img);
      if (!finalUrl) continue;

      out.push({
        url: finalUrl,
        previewUrl: preview && preview !== finalUrl ? preview : undefined,
        contentType: this.contentTypeFor(finalUrl),
        sensitive,
      });
    }

    return out;
  }

  /**
   * Strip Mastodon's ?sensitive=true marker and report whether it was present.
   * The marker is metadata, not part of the media URL, and leaving it in can
   * confuse booru deduplication and content-type sniffing.
   */
  private normalizeUrl(href: string): { url: string; sensitive: boolean } {
    try {
      const u = new URL(href, document.location.href);
      const sensitive = u.searchParams.get("sensitive") === "true";
      u.searchParams.delete("sensitive");
      return { url: u.toString(), sensitive };
    } catch {
      return { url: href, sensitive: false };
    }
  }

  /** Pick the highest-resolution candidate out of an <img srcset>. */
  private largestFromSrcset(img: HTMLImageElement | null): string | undefined {
    if (!img) return undefined;
    const srcset = img.getAttribute("srcset");
    if (!srcset) return img.src || undefined;

    let best: { url: string; w: number } | undefined;
    for (const part of srcset.split(",")) {
      const [urlPart, descriptor] = part.trim().split(/\s+/);
      if (!urlPart) continue;
      const w = descriptor?.endsWith("w") ? parseInt(descriptor, 10) : 0;
      if (!best || w > best.w) best = { url: urlPart, w };
    }
    return best?.url;
  }

  private contentTypeFor(url: string): ContentType {
    const path = url.split("?")[0].toLowerCase();
    if (/\.(mp4|webm|mov|m4v)$/.test(path)) return "video";
    return "image";
  }
}