image-processing/electron/main/image-processor.ts

import { createHash, randomBytes } from 'node:crypto'
import { promises as fs } from 'node:fs'
import path from 'node:path'
import sharp from 'sharp'
import type {
  ImageFileItem,
  ProcessOptions,
  ProcessProgress,
  ProcessSummary,
} from '../../src/shared/types'

export const IMAGE_INPUT_EXTENSIONS = [
  'jpg', 'jpeg', 'jpe', 'jfif',
  'png', 'webp', 'gif', 'tif', 'tiff', 'avif', 'svg',
]

const IMAGE_EXTENSIONS = new Set(IMAGE_INPUT_EXTENSIONS.map((extension) => `.${extension}`))
const SOF_MARKERS = new Set([
  0xc0, 0xc1, 0xc2, 0xc3, 0xc5, 0xc6, 0xc7,
  0xc9, 0xca, 0xcb, 0xcd, 0xce, 0xcf,
])

export class ImageProcessor {
  private cancelled = false

  cancel(): void {
    this.cancelled = true
  }

  async inspect(inputPaths: string[]): Promise<ImageFileItem[]> {
    const files = await collectImagePaths(inputPaths)
    return Promise.all(files.map(async (filePath) => {
      const stat = await fs.stat(filePath)
      try {
        const metadata = await sharp(filePath).metadata()
        return {
          id: fileId(filePath),
          path: filePath,
          name: path.basename(filePath),
          size: stat.size,
          width: metadata.width ?? null,
          height: metadata.height ?? null,
          status: 'pending' as const,
        }
      } catch (error) {
        return {
          id: fileId(filePath),
          path: filePath,
          name: path.basename(filePath),
          size: stat.size,
          width: null,
          height: null,
          status: 'failed' as const,
          error: error instanceof Error ? error.message : '无法读取图片',
        }
      }
    }))
  }

  async process(
    options: ProcessOptions,
    onProgress: (progress: ProcessProgress) => void,
  ): Promise<ProcessSummary> {
    validateOptions(options)
    this.cancelled = false
    let completed = 0
    let failed = 0

    for (let index = 0; index < options.paths.length; index++) {
      const inputPath = options.paths[index]
      if (this.cancelled) {
        break
      }
      const item = await inspectSingle(inputPath)
      item.status = 'processing'
      onProgress(progressFor(index, options.paths.length, item, false))
      try {
        item.outputPath = await processSingleImage(inputPath, options, () => this.cancelled)
        item.status = this.cancelled ? 'cancelled' : 'completed'
        if (item.status === 'completed') {
          completed += 1
        }
      } catch (error) {
        item.status = this.cancelled ? 'cancelled' : 'failed'
        item.error = error instanceof Error ? error.message : '处理失败'
        if (!this.cancelled) {
          failed += 1
        }
      }
      onProgress(progressFor(index + 1, options.paths.length, item, this.cancelled))
    }

    return {
      total: options.paths.length,
      completed,
      failed,
      cancelled: this.cancelled,
    }
  }
}

export async function processSingleImage(
  inputPath: string,
  options: ProcessOptions,
  isCancelled: () => boolean = () => false,
): Promise<string> {
  const metadata = await sharp(inputPath, { page: 0, pages: 1 }).metadata()
  if (!metadata.width || !metadata.height || metadata.height <= 2) {
    throw new Error('图片尺寸过小，无法执行 AI 混淆处理')
  }

  const encoded = await sharp(inputPath, { page: 0, pages: 1 })
    .flatten({ background: '#ffffff' })
    .extract({ left: 0, top: 0, width: metadata.width, height: metadata.height - 2 })
    .jpeg({ quality: 95 })
    .toBuffer()
  const heightRestored = restoreSofHeight(encoded, metadata.height)
  const withoutEoi = removeLastEoi(heightRestored)
  const outputPath = buildOutputPath(inputPath, options)
  const temporaryPath = `${outputPath}.${process.pid}.tmp`

  await fs.mkdir(path.dirname(outputPath), { recursive: true })
  try {
    await fs.writeFile(temporaryPath, withoutEoi)
    let remaining = Math.floor(options.junkSizeMb * 1024 * 1024)
    if (remaining > 0) {
      const handle = await fs.open(temporaryPath, 'a')
      try {
        while (remaining > 0) {
          if (isCancelled()) {
            throw new Error('处理已取消')
          }
          const chunkSize = Math.min(remaining, 1024 * 1024)
          await handle.write(randomBytes(chunkSize))
          remaining -= chunkSize
        }
      } finally {
        await handle.close()
      }
    }
    await fs.rm(outputPath, { force: true })
    await fs.rename(temporaryPath, outputPath)
    return outputPath
  } finally {
    await fs.rm(temporaryPath, { force: true })
  }
}

export function restoreSofHeight(buffer: Buffer, originalHeight: number): Buffer {
  const result = Buffer.from(buffer)
  let index = 0
  while (index < result.length - 1) {
    if (result[index] !== 0xff) {
      index += 1
      continue
    }
    let markerIndex = index + 1
    while (markerIndex < result.length && result[markerIndex] === 0xff) {
      markerIndex += 1
    }
    const marker = result[markerIndex]
    if (SOF_MARKERS.has(marker)) {
      const heightOffset = markerIndex + 4
      if (heightOffset + 1 >= result.length) {
        throw new Error('JPEG SOF 高度字段不完整')
      }
      result[heightOffset] = (originalHeight >> 8) & 0xff
      result[heightOffset + 1] = originalHeight & 0xff
      return result
    }
    if (marker === 0xd8 || marker === 0xd9 || (marker >= 0xd0 && marker <= 0xd7)) {
      index = markerIndex + 1
      continue
    }
    if (markerIndex + 2 >= result.length) {
      break
    }
    const segmentLength = result.readUInt16BE(markerIndex + 1)
    if (segmentLength < 2) {
      break
    }
    index = markerIndex + 1 + segmentLength
  }
  throw new Error('未找到 JPEG SOF 图片尺寸段')
}

export function removeLastEoi(buffer: Buffer): Buffer {
  for (let index = buffer.length - 1; index > 0; index--) {
    if (buffer[index - 1] === 0xff && buffer[index] === 0xd9) {
      return Buffer.concat([buffer.subarray(0, index - 1), buffer.subarray(index + 1)])
    }
  }
  return buffer
}

export function readSofHeight(buffer: Buffer): number | null {
  let index = 0
  while (index < buffer.length - 7) {
    if (buffer[index] === 0xff && SOF_MARKERS.has(buffer[index + 1])) {
      return buffer.readUInt16BE(index + 5)
    }
    index += 1
  }
  return null
}

async function collectImagePaths(inputPaths: string[]): Promise<string[]> {
  const result: string[] = []
  const seen = new Set<string>()
  const stack = [...inputPaths]
  while (stack.length > 0) {
    const candidate = stack.pop()
    if (!candidate) continue
    let stat
    try {
      stat = await fs.lstat(candidate)
    } catch {
      continue
    }
    if (stat.isSymbolicLink()) continue
    if (stat.isDirectory()) {
      const entries = await fs.readdir(candidate, { withFileTypes: true })
      for (const entry of entries) {
        if (!entry.isSymbolicLink()) {
          stack.push(path.join(candidate, entry.name))
        }
      }
    } else if (stat.isFile() && IMAGE_EXTENSIONS.has(path.extname(candidate).toLowerCase())) {
      const absolutePath = path.resolve(candidate)
      if (!seen.has(absolutePath)) {
        seen.add(absolutePath)
        result.push(absolutePath)
      }
    }
  }
  return result.sort((left, right) => left.localeCompare(right, 'zh-CN'))
}

async function inspectSingle(filePath: string): Promise<ImageFileItem> {
  const [stat, metadata] = await Promise.all([
    fs.stat(filePath),
    sharp(filePath).metadata(),
  ])
  return {
    id: fileId(filePath),
    path: filePath,
    name: path.basename(filePath),
    size: stat.size,
    width: metadata.width ?? null,
    height: metadata.height ?? null,
    status: 'pending',
  }
}

function buildOutputPath(inputPath: string, options: ProcessOptions): string {
  const outputDirectory = options.useSourceDirectory
    ? path.dirname(inputPath)
    : path.resolve(options.outputDirectory || '')
  const basename = path.basename(inputPath, path.extname(inputPath))
  const suffix = options.outputFormat === 'file-jpg' ? '.file.jpg' : '.file'
  return path.join(outputDirectory, `${basename}${suffix}`)
}

function validateOptions(options: ProcessOptions): void {
  if (!options.paths.length) {
    throw new Error('请先添加需要处理的图片')
  }
  if (!options.useSourceDirectory && !options.outputDirectory) {
    throw new Error('请选择输出目录')
  }
  if (!Number.isFinite(options.junkSizeMb)
      || options.junkSizeMb < 0 || options.junkSizeMb > 100) {
    throw new Error('追加数据大小必须在 0 到 100 MB 之间')
  }
}

function fileId(filePath: string): string {
  return createHash('sha1').update(filePath).digest('hex')
}

function progressFor(
  current: number,
  total: number,
  item: ImageFileItem,
  cancelled: boolean,
): ProcessProgress {
  return {
    current,
    total,
    percent: total === 0 ? 0 : Math.round((current / total) * 100),
    item,
    cancelled,
  }
}