Image Steganography — Swift source

// steganography — least-significant-bit (LSB) steganography on RGBA pixel
// data, with optional AES-256-GCM encryption.
//
// Language: Swift 5.9+ (Foundation + CryptoKit + CommonCrypto)
// Ported from src/lib/steganography.ts
// display source — part of CosmoDev's polyglot tool pages
//
// Wire format (the "payload" hidden in the pixels):
//   4-byte big-endian header, then the body. The header's top bit is an
//   encryption flag (1 = body is salt+IV+AES-GCM ciphertext, 0 = body is raw
//   UTF-8); the low 31 bits are the body length in bytes. The flag makes the
//   "password required" / "not password-protected" errors deterministic.
//
// Payload bits are written MSB-first, one per R/G/B channel in raster order
// (Alpha is never touched): bit i lands in pixel floor(i/3), channel i%3.
// Capacity = floor(width * height * 3 / 8) payload bytes.
//
// Encryption mirrors the TS reference: PBKDF2-SHA256 (100k iterations,
// 16-byte random salt) derives a 256-bit key via CommonCrypto; CryptoKit's
// AES-GCM encrypts with a 12-byte random IV. Works on any RGBA buffer, so
// tests build synthetic pixels and callers pass CGImage-derived data.

import Foundation
import CryptoKit
import CommonCrypto

let PBKDF2_ITERATIONS = 100_000
let STEG_SALT_BYTES = 16
let STEG_IV_BYTES = 12
let GCM_TAG_BYTES = 16
/// Salt + IV + GCM tag overhead added to the body when a password is used.
let ENCRYPTION_OVERHEAD_BYTES = STEG_SALT_BYTES + STEG_IV_BYTES + GCM_TAG_BYTES
/// The 4-byte length header is also stored in the pixels, so it consumes capacity.
let HEADER_BYTES = 4

// MARK: - Types

/// RGBA pixel data as the browser's ImageData shape: width * height pixels,
/// 4 bytes each (R, G, B, A).
struct StegoImageData {
    let width: Int
    let height: Int
    let data: [UInt8]
}

enum StegoError: Error, LocalizedError {
    case invalidDimensions
    case emptyPassword
    case messageTooLong(bodyBytes: Int, maxBody: Int)
    case noHiddenMessage
    case encryptedNeedsPassword
    case notPasswordProtected
    case decryptionFailed

    var errorDescription: String? {
        switch self {
        case .invalidDimensions:
            return "Width and height must be positive integers"
        case .emptyPassword:
            return "Password must not be empty"
        case .messageTooLong(let body, let max):
            return "Message too long: \(body) bytes with overhead, but this image can hold at most \(max) bytes of message"
        case .noHiddenMessage:
            return "No hidden message found in this image"
        case .encryptedNeedsPassword:
            return "This image contains an encrypted message - a password is required"
        case .notPasswordProtected:
            return "This message is not password-protected - extract without a password"
        case .decryptionFailed:
            return "Decryption failed - wrong password or corrupted data"
        }
    }
}

// MARK: - Capacity

/// Max payload bytes (header + body) an image of this size can carry.
func calculateCapacity(width: Int, height: Int) throws -> Int {
    if width <= 0 || height <= 0 { throw StegoError.invalidDimensions }
    return width * height * 3 / 8
}

// MARK: - Encryption (mirrors src/lib/text-encryptor.ts in the TS build)

/// PBKDF2-SHA256 (100k iterations) -> a 256-bit AES-GCM key.
func deriveKey(password: String, salt: [UInt8]) -> SymmetricKey {
    var keyBytes = [UInt8](repeating: 0, count: 32)
    let pw = Array(password.utf8).map { CChar(bitPattern: $0) }
    _ = CCKeyDerivationPBKDF(
        CCPBKDFAlgorithm(kCCPBKDF2),
        pw, pw.count,
        salt, salt.count,
        CCPseudoRandomAlgorithm(kCCPRFHmacAlgSHA256),
        UInt32(PBKDF2_ITERATIONS),
        &keyBytes, keyBytes.count)
    return SymmetricKey(data: Data(keyBytes))
}

/// 32 CSPRNG bytes — SystemRandomNumberGenerator is the system CSPRNG
/// (SecRandom on Apple platforms), the counterpart of crypto.getRandomValues.
func randomBytes(_ length: Int) -> [UInt8] {
    var csprng = SystemRandomNumberGenerator()
    return (0..<length).map { _ in UInt8.random(in: .min ... .max, using: &csprng) }
}

/// AES-256-GCM encrypt bytes -> packed salt + IV + ciphertext (+ tag).
func encryptBytes(_ plain: [UInt8], password: String) throws -> [UInt8] {
    let salt = randomBytes(STEG_SALT_BYTES)
    let iv = randomBytes(STEG_IV_BYTES)
    let key = deriveKey(password: password, salt: salt)
    let sealed = try AES.GCM.seal(Data(plain), using: key,
                                  nonce: AES.GCM.Nonce(data: Data(iv)))
    // sealed.ciphertext + sealed.tag == WebCrypto's encrypt output shape.
    return salt + iv + sealed.ciphertext + sealed.tag
}

/// Unpack and AES-256-GCM decrypt a salt + IV + ciphertext payload.
func decryptBytes(_ packed: [UInt8], password: String) throws -> [UInt8] {
    let salt = Array(packed.prefix(STEG_SALT_BYTES))
    let iv = Array(packed.dropFirst(STEG_SALT_BYTES).prefix(STEG_IV_BYTES))
    let body = Array(packed.dropFirst(STEG_SALT_BYTES + STEG_IV_BYTES))
    let key = deriveKey(password: password, salt: salt)
    do {
        let box = try AES.GCM.SealedBox(
            nonce: AES.GCM.Nonce(data: Data(iv)),
            ciphertext: Data(body.dropLast(GCM_TAG_BYTES)),
            tag: Data(body.suffix(GCM_TAG_BYTES)))
        return Array(try AES.GCM.open(box, using: key))
    } catch {
        throw StegoError.decryptionFailed
    }
}

// MARK: - Bit embedding

/// Write `payload` into the LSBs of the R/G/B channels; returns copied pixels.
func embedBits(_ data: [UInt8], payload: [UInt8]) -> [UInt8] {
    var out = data  // copy - the input is never mutated
    let totalBits = payload.count * 8
    for i in 0..<totalBits {
        let byte = payload[i >> 3]
        let bit = (byte >> UInt8(7 - (i & 7))) & 1
        let px = i / 3
        let channel = i % 3
        let idx = px * 4 + channel
        out[idx] = (out[idx] & 0xfe) | bit
    }
    return out
}

/// Read `count` payload bytes back out of the R/G/B LSBs.
func extractBits(_ data: [UInt8], offsetBytes: Int, count: Int) -> [UInt8] {
    var out = [UInt8](repeating: 0, count: count)
    let startBit = offsetBytes * 8
    for i in 0..<(count * 8) {
        let bitIndex = startBit + i
        let px = bitIndex / 3
        let channel = bitIndex % 3
        let bit = data[px * 4 + channel] & 1
        out[i >> 3] |= bit << UInt8(7 - (i & 7))
    }
    return out
}

/// Big-endian UInt32 from a 4-byte buffer.
func readHeader(_ bytes: [UInt8]) -> UInt32 {
    (UInt32(bytes[0]) << 24) | (UInt32(bytes[1]) << 16)
        | (UInt32(bytes[2]) << 8) | UInt32(bytes[3])
}

// MARK: - Public API

/// Hide `message` inside a copy of `imageData`'s pixels (LSB of R/G/B) and
/// return the modified pixel data. With `password`, the message body is
/// AES-256-GCM encrypted first. Throws if the message (including header and
/// encryption overhead) exceeds the image capacity, or on an empty password.
func hideMessage(_ imageData: StegoImageData, message: String, password: String? = nil) throws -> StegoImageData {
    if let password = password, password.isEmpty {
        throw StegoError.emptyPassword
    }
    let capacity = try calculateCapacity(width: imageData.width, height: imageData.height)
    let plain = Array(message.utf8)
    let body = try password != nil ? encryptBytes(plain, password: password!) : plain
    var payload = [UInt8](repeating: 0, count: HEADER_BYTES + body.count)
    // Top bit = encryption flag; low 31 bits = body length. Big-endian.
    let header = UInt32(body.count) | (password != nil ? 0x8000_0000 : 0)
    payload[0] = UInt8((header >> 24) & 0xff)
    payload[1] = UInt8((header >> 16) & 0xff)
    payload[2] = UInt8((header >> 8) & 0xff)
    payload[3] = UInt8(header & 0xff)
    payload.replaceSubrange(HEADER_BYTES..., with: body)
    if payload.count > capacity {
        throw StegoError.messageTooLong(bodyBytes: body.count, maxBody: capacity - HEADER_BYTES)
    }
    return StegoImageData(
        width: imageData.width,
        height: imageData.height,
        data: embedBits(imageData.data, payload: payload))
}

/// Read the hidden message out of `imageData`'s pixels. Throws when the pixels
/// carry no valid payload, when the payload is encrypted but no password is
/// given, when a password is given but the payload is plaintext, and on a
/// wrong password (GCM authentication failure).
func extractMessage(_ imageData: StegoImageData, password: String? = nil) throws -> String {
    if let password = password, password.isEmpty {
        throw StegoError.emptyPassword
    }
    let capacity = try calculateCapacity(width: imageData.width, height: imageData.height)
    let header = readHeader(extractBits(imageData.data, offsetBytes: 0, count: HEADER_BYTES))
    let encrypted = (header & 0x8000_0000) != 0
    let length = Int(header & 0x7fff_ffff)
    if length == 0 && !encrypted { return "" }
    let minLength = encrypted ? STEG_SALT_BYTES + STEG_IV_BYTES + GCM_TAG_BYTES : 1
    if HEADER_BYTES + length > capacity || length < minLength {
        throw StegoError.noHiddenMessage
    }
    let body = extractBits(imageData.data, offsetBytes: HEADER_BYTES, count: length)
    if !encrypted {
        if password != nil {
            throw StegoError.notPasswordProtected
        }
        // Strict UTF-8 decode — invalid bytes mean "not a real payload",
        // matching the TS TextDecoder(fatal: true) behavior.
        guard let text = String(bytes: body, encoding: .utf8) else {
            throw StegoError.noHiddenMessage
        }
        return text
    }
    guard let password = password else {
        throw StegoError.encryptedNeedsPassword
    }
    let plain = try decryptBytes(body, password: password)
    return String(decoding: plain, as: UTF8.self)
}