なんかもうめっちゃ変えた

1 files changed, 382 insertions, 0 deletions

diff --git a/packages/backend/src/core/FileInfoService.ts b/packages/backend/src/core/FileInfoService.ts
new file mode 100644
index 0000000000..45c94dafdc
--- /dev/null
+++ b/packages/backend/src/core/FileInfoService.ts
@@ -0,0 +1,382 @@
+import * as fs from 'node:fs';
+import * as crypto from 'node:crypto';
+import { join } from 'node:path';
+import * as stream from 'node:stream';
+import * as util from 'node:util';
+import { Inject, Injectable } from '@nestjs/common';
+import { FSWatcher } from 'chokidar';
+import { fileTypeFromFile } from 'file-type';
+import FFmpeg from 'fluent-ffmpeg';
+import isSvg from 'is-svg';
+import probeImageSize from 'probe-image-size';
+import { type predictionType } from 'nsfwjs';
+import sharp from 'sharp';
+import { encode } from 'blurhash';
+import { createTempDir } from '@/misc/create-temp.js';
+import { AiService } from '@/core/AiService.js';
+
+const pipeline = util.promisify(stream.pipeline);
+
+export type FileInfo = {
+	size: number;
+	md5: string;
+	type: {
+		mime: string;
+		ext: string | null;
+	};
+	width?: number;
+	height?: number;
+	orientation?: number;
+	blurhash?: string;
+	sensitive: boolean;
+	porn: boolean;
+	warnings: string[];
+};
+
+const TYPE_OCTET_STREAM = {
+	mime: 'application/octet-stream',
+	ext: null,
+};
+
+const TYPE_SVG = {
+	mime: 'image/svg+xml',
+	ext: 'svg',
+};
+@Injectable()
+export class FileInfoService {
+	constructor(
+		private aiService: AiService,
+	) {
+	}
+
+	/**
+	 * Get file information
+	 */
+	public async getFileInfo(path: string, opts: {
+		skipSensitiveDetection: boolean;
+		sensitiveThreshold?: number;
+		sensitiveThresholdForPorn?: number;
+		enableSensitiveMediaDetectionForVideos?: boolean;
+	}): Promise<FileInfo> {
+		const warnings = [] as string[];
+
+		const size = await this.getFileSize(path);
+		const md5 = await this.#calcHash(path);
+
+		let type = await this.detectType(path);
+
+		// image dimensions
+		let width: number | undefined;
+		let height: number | undefined;
+		let orientation: number | undefined;
+
+		if (['image/jpeg', 'image/gif', 'image/png', 'image/apng', 'image/webp', 'image/bmp', 'image/tiff', 'image/svg+xml', 'image/vnd.adobe.photoshop'].includes(type.mime)) {
+			const imageSize = await this.#detectImageSize(path).catch(e => {
+				warnings.push(`detectImageSize failed: ${e}`);
+				return undefined;
+			});
+
+			// うまく判定できない画像は octet-stream にする
+			if (!imageSize) {
+				warnings.push('cannot detect image dimensions');
+				type = TYPE_OCTET_STREAM;
+			} else if (imageSize.wUnits === 'px') {
+				width = imageSize.width;
+				height = imageSize.height;
+				orientation = imageSize.orientation;
+
+				// 制限を超えている画像は octet-stream にする
+				if (imageSize.width > 16383 || imageSize.height > 16383) {
+					warnings.push('image dimensions exceeds limits');
+					type = TYPE_OCTET_STREAM;
+				}
+			} else {
+				warnings.push(`unsupported unit type: ${imageSize.wUnits}`);
+			}
+		}
+
+		let blurhash: string | undefined;
+
+		if (['image/jpeg', 'image/gif', 'image/png', 'image/apng', 'image/webp', 'image/svg+xml'].includes(type.mime)) {
+			blurhash = await this.#getBlurhash(path).catch(e => {
+				warnings.push(`getBlurhash failed: ${e}`);
+				return undefined;
+			});
+		}
+
+		let sensitive = false;
+		let porn = false;
+
+		if (!opts.skipSensitiveDetection) {
+			await this.#detectSensitivity(
+				path,
+				type.mime,
+				opts.sensitiveThreshold ?? 0.5,
+				opts.sensitiveThresholdForPorn ?? 0.75,
+				opts.enableSensitiveMediaDetectionForVideos ?? false,
+			).then(value => {
+				[sensitive, porn] = value;
+			}, error => {
+				warnings.push(`detectSensitivity failed: ${error}`);
+			});
+		}
+
+		return {
+			size,
+			md5,
+			type,
+			width,
+			height,
+			orientation,
+			blurhash,
+			sensitive,
+			porn,
+			warnings,
+		};
+	}
+
+	async #detectSensitivity(source: string, mime: string, sensitiveThreshold: number, sensitiveThresholdForPorn: number, analyzeVideo: boolean): Promise<[sensitive: boolean, porn: boolean]> {
+		let sensitive = false;
+		let porn = false;
+	
+		function judgePrediction(result: readonly predictionType[]): [sensitive: boolean, porn: boolean] {
+			let sensitive = false;
+			let porn = false;
+	
+			if ((result.find(x => x.className === 'Sexy')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
+			if ((result.find(x => x.className === 'Hentai')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
+			if ((result.find(x => x.className === 'Porn')?.probability ?? 0) > sensitiveThreshold) sensitive = true;
+	
+			if ((result.find(x => x.className === 'Porn')?.probability ?? 0) > sensitiveThresholdForPorn) porn = true;
+	
+			return [sensitive, porn];
+		}
+	
+		if (['image/jpeg', 'image/png', 'image/webp'].includes(mime)) {
+			const result = await this.aiService.detectSensitive(source);
+			if (result) {
+				[sensitive, porn] = judgePrediction(result);
+			}
+		} else if (analyzeVideo && (mime === 'image/apng' || mime.startsWith('video/'))) {
+			const [outDir, disposeOutDir] = await createTempDir();
+			try {
+				const command = FFmpeg()
+					.input(source)
+					.inputOptions([
+						'-skip_frame', 'nokey', // 可能ならキーフレームのみを取得してほしいとする（そうなるとは限らない）
+						'-lowres', '3', // 元の画質でデコードする必要はないので 1/8 画質でデコードしてもよいとする（そうなるとは限らない）
+					])
+					.noAudio()
+					.videoFilters([
+						{
+							filter: 'select', // フレームのフィルタリング
+							options: {
+								e: 'eq(pict_type,PICT_TYPE_I)', // I-Frame のみをフィルタする（VP9 とかはデコードしてみないとわからないっぽい）
+							},
+						},
+						{
+							filter: 'blackframe', // 暗いフレームの検出
+							options: {
+								amount: '0', // 暗さに関わらず全てのフレームで測定値を取る
+							},
+						},
+						{
+							filter: 'metadata',
+							options: {
+								mode: 'select', // フレーム選択モード
+								key: 'lavfi.blackframe.pblack', // フレームにおける暗部の百分率（前のフィルタからのメタデータを参照する）
+								value: '50',
+								function: 'less', // 50% 未満のフレームを選択する（50% 以上暗部があるフレームだと誤検知を招くかもしれないので）
+							},
+						},
+						{
+							filter: 'scale',
+							options: {
+								w: 299,
+								h: 299,
+							},
+						},
+					])
+					.format('image2')
+					.output(join(outDir, '%d.png'))
+					.outputOptions(['-vsync', '0']); // 可変フレームレートにすることで穴埋めをさせない
+				const results: ReturnType<typeof judgePrediction>[] = [];
+				let frameIndex = 0;
+				let targetIndex = 0;
+				let nextIndex = 1;
+				for await (const path of this.#asyncIterateFrames(outDir, command)) {
+					try {
+						const index = frameIndex++;
+						if (index !== targetIndex) {
+							continue;
+						}
+						targetIndex = nextIndex;
+						nextIndex += index; // fibonacci sequence によってフレーム数制限を掛ける
+						const result = await this.aiService.detectSensitive(path);
+						if (result) {
+							results.push(judgePrediction(result));
+						}
+					} finally {
+						fs.promises.unlink(path);
+					}
+				}
+				sensitive = results.filter(x => x[0]).length >= Math.ceil(results.length * sensitiveThreshold);
+				porn = results.filter(x => x[1]).length >= Math.ceil(results.length * sensitiveThresholdForPorn);
+			} finally {
+				disposeOutDir();
+			}
+		}
+	
+		return [sensitive, porn];
+	}
+	
+	async *#asyncIterateFrames(cwd: string, command: FFmpeg.FfmpegCommand): AsyncGenerator<string, void> {
+		const watcher = new FSWatcher({
+			cwd,
+			disableGlobbing: true,
+		});
+		let finished = false;
+		command.once('end', () => {
+			finished = true;
+			watcher.close();
+		});
+		command.run();
+		for (let i = 1; true; i++) { // eslint-disable-line @typescript-eslint/no-unnecessary-condition
+			const current = `${i}.png`;
+			const next = `${i + 1}.png`;
+			const framePath = join(cwd, current);
+			if (await this.#exists(join(cwd, next))) {
+				yield framePath;
+			} else if (!finished) { // eslint-disable-line @typescript-eslint/no-unnecessary-condition
+				watcher.add(next);
+				await new Promise<void>((resolve, reject) => {
+					watcher.on('add', function onAdd(path) {
+						if (path === next) { // 次フレームの書き出しが始まっているなら、現在フレームの書き出しは終わっている
+							watcher.unwatch(current);
+							watcher.off('add', onAdd);
+							resolve();
+						}
+					});
+					command.once('end', resolve); // 全てのフレームを処理し終わったなら、最終フレームである現在フレームの書き出しは終わっている
+					command.once('error', reject);
+				});
+				yield framePath;
+			} else if (await this.#exists(framePath)) {
+				yield framePath;
+			} else {
+				return;
+			}
+		}
+	}
+	
+	#exists(path: string): Promise<boolean> {
+		return fs.promises.access(path).then(() => true, () => false);
+	}
+
+	/**
+	 * Detect MIME Type and extension
+	 */
+	public async detectType(path: string): Promise<{
+	mime: string;
+	ext: string | null;
+}> {
+	// Check 0 byte
+		const fileSize = await this.getFileSize(path);
+		if (fileSize === 0) {
+			return TYPE_OCTET_STREAM;
+		}
+
+		const type = await fileTypeFromFile(path);
+
+		if (type) {
+		// XMLはSVGかもしれない
+			if (type.mime === 'application/xml' && await this.checkSvg(path)) {
+				return TYPE_SVG;
+			}
+
+			return {
+				mime: type.mime,
+				ext: type.ext,
+			};
+		}
+
+		// 種類が不明でもSVGかもしれない
+		if (await this.checkSvg(path)) {
+			return TYPE_SVG;
+		}
+
+		// それでも種類が不明なら application/octet-stream にする
+		return TYPE_OCTET_STREAM;
+	}
+
+	/**
+	 * Check the file is SVG or not
+	 */
+	public async checkSvg(path: string) {
+		try {
+			const size = await this.getFileSize(path);
+			if (size > 1 * 1024 * 1024) return false;
+			return isSvg(fs.readFileSync(path));
+		} catch {
+			return false;
+		}
+	}
+
+	/**
+	 * Get file size
+	 */
+	public async getFileSize(path: string): Promise<number> {
+		const getStat = util.promisify(fs.stat);
+		return (await getStat(path)).size;
+	}
+
+	/**
+	 * Calculate MD5 hash
+	 */
+	async #calcHash(path: string): Promise<string> {
+		const hash = crypto.createHash('md5').setEncoding('hex');
+		await pipeline(fs.createReadStream(path), hash);
+		return hash.read();
+	}
+
+	/**
+	 * Detect dimensions of image
+	 */
+	async #detectImageSize(path: string): Promise<{
+	width: number;
+	height: number;
+	wUnits: string;
+	hUnits: string;
+	orientation?: number;
+}> {
+		const readable = fs.createReadStream(path);
+		const imageSize = await probeImageSize(readable);
+		readable.destroy();
+		return imageSize;
+	}
+
+	/**
+	 * Calculate average color of image
+	 */
+	#getBlurhash(path: string): Promise<string> {
+		return new Promise((resolve, reject) => {
+			sharp(path)
+				.raw()
+				.ensureAlpha()
+				.resize(64, 64, { fit: 'inside' })
+				.toBuffer((err, buffer, { width, height }) => {
+					if (err) return reject(err);
+
+					let hash;
+
+					try {
+						hash = encode(new Uint8ClampedArray(buffer), width, height, 7, 7);
+					} catch (e) {
+						return reject(e);
+					}
+
+					resolve(hash);
+				});
+		});
+	}
+}