Image Steganography — Kotlin source

// Steganography — least-significant-bit (LSB) message hiding on RGBA pixel
// data, with optional AES-256-GCM encryption.
//
// Language: Kotlin 1.9+ (JVM), javax.crypto only.
// Ported from src/lib/steganography.ts — display source, part of CosmoDev's
// polyglot tool pages. Functionally equivalent to the TS reference (which
// drives Web Crypto): the TS side is async because SubtleCrypto is; the JVM's
// javax.crypto calls are synchronous, so this port is too. Works on any
// width/height/RGBA byte triple — synthetic buffers in tests, a Bitmap's
// pixels on Android, or BufferedImage raster data on the desktop.
//
// 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 text-encryptor tool: PBKDF2-SHA256 (100k iterations,
// 16-byte random salt) derives an AES-256 key; GCM encrypts with a 12-byte
// random IV.

import javax.crypto.Cipher
import javax.crypto.spec.GCMParameterSpec
import javax.crypto.spec.SecretKeySpec
import java.nio.charset.CodingErrorAction
import java.security.SecureRandom

private const val PBKDF2_ITERATIONS = 100_000
private const val SALT_BYTES = 16
private const val IV_BYTES = 12
private const val GCM_TAG_BYTES = 16

/** Salt + IV + GCM tag overhead added to the body when a password is used. */
const val ENCRYPTION_OVERHEAD_BYTES = SALT_BYTES + IV_BYTES + GCM_TAG_BYTES

/** The 4-byte length header is also stored in the pixels, so it consumes capacity. */
const val HEADER_BYTES = 4

private const val TAG_BITS = GCM_TAG_BYTES * 8

/** Pixel buffer shape: width x height of RGBA quadruples. */
data class StegoImageData(
    val width: Int,
    val height: Int,
    val data: ByteArray,
) {
    override fun equals(other: Any?): Boolean =
        other is StegoImageData && other.width == width && other.height == height &&
            other.data.contentEquals(data)

    override fun hashCode(): Int = 31 * 31 * width + height * 31 + data.contentHashCode()
}

/** Max payload bytes (header + body) an image of this size can carry. */
fun calculateCapacity(width: Int, height: Int): Int {
    if (width <= 0 || height <= 0) {
        throw IllegalArgumentException("Width and height must be positive integers")
    }
    return width * height * 3 / 8
}

private val RANDOM = SecureRandom()

/** PBKDF2-SHA256 (100k iterations) -> AES-256 key. */
private fun deriveKey(password: String, salt: ByteArray): SecretKeySpec {
    val factory = javax.crypto.SecretKeyFactory.getInstance("PBKDF2WithHmacSHA256")
    val spec = javax.crypto.spec.PBEKeySpec(password.toCharArray(), salt, PBKDF2_ITERATIONS, 256)
    return SecretKeySpec(factory.generateSecret(spec).encoded, "AES")
}

/** AES-256-GCM encrypt bytes -> packed salt + IV + ciphertext (+ tag). */
private fun encryptBytes(plain: ByteArray, password: String): ByteArray {
    val salt = ByteArray(SALT_BYTES).also(RANDOM::nextBytes)
    val iv = ByteArray(IV_BYTES).also(RANDOM::nextBytes)
    val cipher = Cipher.getInstance("AES/GCM/NoPadding")
    cipher.init(Cipher.ENCRYPT_MODE, deriveKey(password, salt), GCMParameterSpec(TAG_BITS, iv))
    val encrypted = cipher.doFinal(plain)
    return salt + iv + encrypted
}

/** Unpack and AES-256-GCM decrypt a salt + IV + ciphertext payload. */
private fun decryptBytes(packed: ByteArray, password: String): ByteArray {
    val salt = packed.copyOfRange(0, SALT_BYTES)
    val iv = packed.copyOfRange(SALT_BYTES, SALT_BYTES + IV_BYTES)
    val data = packed.copyOfRange(SALT_BYTES + IV_BYTES, packed.size)
    return try {
        val cipher = Cipher.getInstance("AES/GCM/NoPadding")
        cipher.init(Cipher.DECRYPT_MODE, deriveKey(password, salt), GCMParameterSpec(TAG_BITS, iv))
        cipher.doFinal(data)
    } catch (_: Exception) {
        throw IllegalArgumentException("Decryption failed - wrong password or corrupted data")
    }
}

/** Strict UTF-8 decode: invalid sequences throw (mirrors TextDecoder fatal:true). */
private fun decodeUtf8Strict(bytes: ByteArray): String {
    val decoder = Charsets.UTF_8.newDecoder()
        .onMalformedInput(CodingErrorAction.REPORT)
        .onUnmappableCharacter(CodingErrorAction.REPORT)
    return decoder.decode(java.nio.ByteBuffer.wrap(bytes)).toString()
}

/** Write `payload` into the LSBs of the R/G/B channels; returns copied pixels. */
private fun embedBits(data: ByteArray, payload: ByteArray): ByteArray {
    val out = data.copyOf() // copy - the input is never mutated
    val totalBits = payload.size * 8
    for (i in 0 until totalBits) {
        val bit = (payload[i shr 3].toInt() shr (7 - (i and 7))) and 1
        val px = i / 3
        val channel = i % 3
        val idx = px * 4 + channel
        out[idx] = ((out[idx].toInt() and 0xFE) or bit).toByte()
    }
    return out
}

/** Read `count` payload bytes back out of the R/G/B LSBs. */
private fun extractBits(data: ByteArray, offsetBytes: Int, count: Int): ByteArray {
    val out = ByteArray(count)
    val startBit = offsetBytes * 8
    for (i in 0 until count * 8) {
        val bitIndex = startBit + i
        val px = bitIndex / 3
        val channel = bitIndex % 3
        val bit = data[px * 4 + channel].toInt() and 1
        out[i shr 3] = (out[i shr 3] + (bit shl (7 - (i and 7)))).toByte()
    }
    return out
}

/** Big-endian u32 -> 4 bytes. */
private fun u32Bytes(n: Int): ByteArray = byteArrayOf(
    (n ushr 24).toByte(), (n ushr 16).toByte(), (n ushr 8).toByte(), n.toByte(),
)

private fun readU32(bytes: ByteArray): Int {
    var n = 0
    for (b in bytes) n = (n shl 8) or (b.toInt() and 0xFF)
    return n
}

/**
 * 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.
 */
fun hideMessage(imageData: StegoImageData, message: String, password: String? = null): StegoImageData {
    if (password != null && password.isEmpty()) {
        throw IllegalArgumentException("Password must not be empty")
    }
    val capacity = calculateCapacity(imageData.width, imageData.height)
    val plain = message.toByteArray(Charsets.UTF_8)
    val body = if (password != null) encryptBytes(plain, password) else plain
    val header = body.size or (if (password != null) 0x8000_0000.toInt() else 0)
    val payload = u32Bytes(header) + body
    if (payload.size > capacity) {
        val maxBody = capacity - HEADER_BYTES
        throw IllegalArgumentException(
            "Message too long: ${body.size} bytes with overhead, but this image can hold at most $maxBody bytes of message"
        )
    }
    return StegoImageData(imageData.width, imageData.height, embedBits(imageData.data, payload))
}

/**
 * Read the hidden message out of `imageData`'s pixels. Throws when the pixels
 * carry no valid payload ("No hidden message found"), 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).
 */
fun extractMessage(imageData: StegoImageData, password: String? = null): String {
    if (password != null && password.isEmpty()) {
        throw IllegalArgumentException("Password must not be empty")
    }
    val capacity = calculateCapacity(imageData.width, imageData.height)
    val header = readU32(extractBits(imageData.data, 0, HEADER_BYTES))
    val encrypted = (header and 0x8000_0000.toInt()) != 0
    val length = header and 0x7fff_ffff
    if (length == 0 && !encrypted) return ""
    if (HEADER_BYTES + length > capacity ||
        length < if (encrypted) SALT_BYTES + IV_BYTES + GCM_TAG_BYTES else 1
    ) {
        throw IllegalArgumentException("No hidden message found in this image")
    }
    val body = extractBits(imageData.data, HEADER_BYTES, length)
    if (!encrypted) {
        if (password != null) {
            throw IllegalArgumentException("This message is not password-protected - extract without a password")
        }
        return try {
            decodeUtf8Strict(body)
        } catch (_: Exception) {
            throw IllegalArgumentException("No hidden message found in this image")
        }
    }
    if (password == null) {
        throw IllegalArgumentException("This image contains an encrypted message - a password is required")
    }
    val plain = decryptBytes(body, password)
    return String(plain, Charsets.UTF_8)
}