Image Steganography — Go source

// Package steganography is the Go twin of CosmoDev's src/lib/steganography.ts
// (dual source: the web lib is TypeScript, the CLI lib is Go — kept in
// lock-step). Pure, never panics; returns errors or zero values. The tests in
// steganography_test.go share vectors with src/lib/steganography.test.ts so
// the two implementations are held to the same contract.
//
// Wire format (mirrors the TS lib exactly):
//
//	4-byte big-endian header, then the body. The header's top bit is an
//	encryption flag (1 = body is salt+IV+AES-256-GCM ciphertext, 0 = body is
//	raw UTF-8); the low 31 bits are the body length in bytes.
//
// Payload bits are written MSB-first, one per R/G/B channel in raster order
// (Alpha is never touched): bit i lands in pixel i/3, channel i%3.
// Capacity = width*height*3/8 payload bytes, rounded down to whole bytes.
//
// Like the TS lib (which works on ImageData-shaped objects, not Image
// objects), this package operates on RGBA byte buffers — image decode/encode
// plumbing stays out.
package steganography

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/pbkdf2"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"fmt"
	"unicode/utf8"
)

const (
	// PBKDF2Iterations is the PBKDF2-SHA256 iteration count for deriving the
	// AES-256 key from a password. Mirrors PBKDF2_ITERATIONS in the TS lib.
	PBKDF2Iterations = 100_000
	// SaltBytes is the random salt length prepended to an encrypted body.
	SaltBytes = 16
	// IVBytes is the random GCM nonce length (12 is the GCM standard size).
	IVBytes = 12
	// GCMTagBytes is the AES-GCM authentication tag length.
	GCMTagBytes = 16
	// EncryptionOverheadBytes is the salt+IV+GCM tag overhead added to the
	// body when a password is used. Mirrors ENCRYPTION_OVERHEAD_BYTES.
	EncryptionOverheadBytes = SaltBytes + IVBytes + GCMTagBytes
	// HeaderBytes is the 4-byte length header's size; it is also stored in
	// the pixels, so it consumes capacity. Mirrors HEADER_BYTES.
	HeaderBytes = 4

	encryptedFlag uint32 = 0x8000_0000
	lengthMask    uint32 = 0x7fff_ffff
)

// Errors mirror the TS lib's thrown messages; both test suites match on this
// text (the TS one by substring, this one by errors.Is / substring).
var (
	// ErrInvalidDimensions: width and height must be positive integers. (The
	// TS lib also rejects non-integers — unrepresentable as a Go int.)
	ErrInvalidDimensions = errors.New("width and height must be positive integers")
	// ErrEmptyPassword: an empty (but present) password was supplied.
	ErrEmptyPassword = errors.New("password must not be empty")
	// ErrNoMessage: the pixels carry no valid payload.
	ErrNoMessage = errors.New("no hidden message found in this image")
	// ErrPasswordRequired: the payload is encrypted but no password was given.
	ErrPasswordRequired = errors.New("this image contains an encrypted message - a password is required")
	// ErrNotPasswordProtected: a password was given but the payload is plaintext.
	ErrNotPasswordProtected = errors.New("this message is not password-protected - extract without a password")
	// ErrDecryptionFailed: GCM authentication failed (wrong password or
	// corrupted data). Mirrors the TS catch-and-rethrow in decryptBytes.
	ErrDecryptionFailed = errors.New("decryption failed - wrong password or corrupted data")
	// ErrBadPixelData: Data is not exactly width*height*4 bytes. Go-specific
	// guard: the TS lib cannot panic on a short buffer (typed-array indexing
	// silently drops out-of-range reads/writes), but this port must not panic.
	ErrBadPixelData = errors.New("pixel data must be exactly width*height*4 bytes")
)

// ImageData is the minimal structural type satisfied by the browser's
// ImageData. Mirrors StegoImageData in the TS lib.
type ImageData struct {
	Width  int
	Height int
	Data   []byte // RGBA, one byte per channel, len == Width*Height*4
}

// CalculateCapacity returns the max payload bytes (header + body) an image of
// this size can carry. Mirrors calculateCapacity in the TS lib.
func CalculateCapacity(width, height int) (int, error) {
	if width <= 0 || height <= 0 {
		return 0, ErrInvalidDimensions
	}
	return width * height * 3 / 8, nil
}

// validate checks the pixel-buffer shape so the bit loops below can never
// index out of range.
func (img ImageData) validate() error {
	if len(img.Data) != img.Width*img.Height*4 {
		return fmt.Errorf("%w: got %d bytes for %dx%d", ErrBadPixelData, len(img.Data), img.Width, img.Height)
	}
	return nil
}

// newGCM derives the AES-256 key (PBKDF2-SHA256, 100k iterations) from
// password+salt and returns a GCM AEAD. Mirrors deriveKey() in the TS lib —
// the derived bytes are identical, so payloads interoperate across TS/Go.
func newGCM(password string, salt []byte) (cipher.AEAD, error) {
	key, err := pbkdf2.Key(sha256.New, password, salt, PBKDF2Iterations, 32)
	if err != nil {
		return nil, fmt.Errorf("deriving key: %w", err)
	}
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, fmt.Errorf("creating cipher: %w", err)
	}
	return cipher.NewGCM(block)
}

// encryptBytes AES-256-GCM encrypts plain and returns packed
// salt+IV+ciphertext(+tag). Mirrors encryptBytes() in the TS lib: Seal
// appends the 16-byte tag, exactly like WebCrypto's ciphertext.
func encryptBytes(plain []byte, password string) ([]byte, error) {
	salt := make([]byte, SaltBytes)
	iv := make([]byte, IVBytes)
	if _, err := rand.Read(salt); err != nil {
		return nil, fmt.Errorf("reading random salt: %w", err)
	}
	if _, err := rand.Read(iv); err != nil {
		return nil, fmt.Errorf("reading random IV: %w", err)
	}
	gcm, err := newGCM(password, salt)
	if err != nil {
		return nil, err
	}
	packed := make([]byte, 0, SaltBytes+IVBytes+len(plain)+GCMTagBytes)
	packed = append(packed, salt...)
	packed = append(packed, iv...)
	return gcm.Seal(packed, iv, plain, nil), nil
}

// decryptBytes unpacks and AES-256-GCM decrypts a salt+IV+ciphertext payload.
// Mirrors decryptBytes() in the TS lib, including the wrong-password failure.
func decryptBytes(packed []byte, password string) ([]byte, error) {
	if len(packed) < SaltBytes+IVBytes+GCMTagBytes {
		return nil, ErrDecryptionFailed
	}
	salt := packed[:SaltBytes]
	iv := packed[SaltBytes : SaltBytes+IVBytes]
	data := packed[SaltBytes+IVBytes:]
	gcm, err := newGCM(password, salt)
	if err != nil {
		return nil, err
	}
	plain, err := gcm.Open(nil, iv, data, nil)
	if err != nil {
		return nil, ErrDecryptionFailed
	}
	return plain, nil
}

// embedBits writes payload into the LSBs of the R/G/B channels of a copy of
// data and returns it. Mirrors embedBits() in the TS lib — the input is never
// mutated.
func embedBits(data, payload []byte) []byte {
	out := make([]byte, len(data))
	copy(out, data)
	for i := 0; i < len(payload)*8; i++ {
		bit := (payload[i>>3] >> (7 - i&7)) & 1
		idx := i/3*4 + i%3
		out[idx] = out[idx]&0xfe | bit
	}
	return out
}

// extractBits reads count payload bytes back out of the R/G/B LSBs, starting
// at offsetBytes into the payload stream. Mirrors extractBits() in the TS lib.
func extractBits(data []byte, offsetBytes, count int) []byte {
	out := make([]byte, count)
	startBit := offsetBytes * 8
	for i := 0; i < count*8; i++ {
		bitIndex := startBit + i
		bit := data[bitIndex/3*4+bitIndex%3] & 1
		out[i>>3] |= bit << (7 - i&7)
	}
	return out
}

// HideMessage hides message inside a copy of img's pixels (LSB of R/G/B) and
// returns the modified image. With a non-nil non-empty password the message
// body is AES-256-GCM encrypted first.
//
// The password is a *string so nil means "no password" (plaintext) while a
// non-nil empty string is an error — mirroring the TS lib's
// undefined/''/value tristate (the same pattern slugify's Options.Separator
// uses). It fails with ErrEmptyPassword, ErrInvalidDimensions /
// ErrBadPixelData, or a "Message too long" error when the payload (header +
// body + encryption overhead) exceeds the image capacity.
func HideMessage(img ImageData, message string, password *string) (ImageData, error) {
	if password != nil && *password == "" {
		return ImageData{}, ErrEmptyPassword
	}
	if err := img.validate(); err != nil {
		return ImageData{}, err
	}
	capacity, err := CalculateCapacity(img.Width, img.Height)
	if err != nil {
		return ImageData{}, err
	}
	body := []byte(message)
	if password != nil {
		if body, err = encryptBytes(body, *password); err != nil {
			return ImageData{}, err
		}
	}
	header := uint32(len(body))
	if password != nil {
		header |= encryptedFlag
	}
	payload := make([]byte, HeaderBytes+len(body))
	binary.BigEndian.PutUint32(payload[:HeaderBytes], header)
	copy(payload[HeaderBytes:], body)
	if len(payload) > capacity {
		return ImageData{}, fmt.Errorf(
			"Message too long: %d bytes with overhead, but this image can hold at most %d bytes of message",
			len(body), capacity-HeaderBytes)
	}
	return ImageData{
		Width:  img.Width,
		Height: img.Height,
		Data:   embedBits(img.Data, payload),
	}, nil
}

// ExtractMessage reads the hidden message out of img's pixels. It fails with
// ErrNoMessage when the pixels carry no valid payload (including a corrupted
// one that decodes to invalid UTF-8), ErrPasswordRequired when the payload is
// encrypted but no password is given, ErrNotPasswordProtected when a password
// is given but the payload is plaintext, and ErrDecryptionFailed on a wrong
// password (GCM authentication failure).
func ExtractMessage(img ImageData, password *string) (string, error) {
	if password != nil && *password == "" {
		return "", ErrEmptyPassword
	}
	if err := img.validate(); err != nil {
		return "", err
	}
	capacity, err := CalculateCapacity(img.Width, img.Height)
	if err != nil {
		return "", err
	}
	header := binary.BigEndian.Uint32(extractBits(img.Data, 0, HeaderBytes))
	encrypted := header&encryptedFlag != 0
	length := header & lengthMask
	if length == 0 && !encrypted {
		return "", nil
	}
	minLength := 1
	if encrypted {
		minLength = SaltBytes + IVBytes + GCMTagBytes
	}
	if int64(HeaderBytes)+int64(length) > int64(capacity) || int64(length) < int64(minLength) {
		return "", ErrNoMessage
	}
	body := extractBits(img.Data, HeaderBytes, int(length))
	if !encrypted {
		if password != nil {
			return "", ErrNotPasswordProtected
		}
		if !utf8.Valid(body) {
			return "", ErrNoMessage
		}
		return string(body), nil
	}
	if password == nil {
		return "", ErrPasswordRequired
	}
	plain, err := decryptBytes(body, *password)
	if err != nil {
		return "", err
	}
	return string(plain), nil
}