// Copyright (C) 2018 Marius Schellenberger

package main

import (
	"crypto/cipher"
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"image"
	"image/color"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"io"
	"math"
	"os"

	"golang.org/x/crypto/chacha20"
	"golang.org/x/crypto/chacha20poly1305"
	"golang.org/x/crypto/hkdf"
)

const (
	KeySize   = chacha20poly1305.KeySize
	NonceSize = chacha20poly1305.NonceSize
	Overhead  = chacha20poly1305.Overhead

	ChunkSize = 512
	FullSize  = ChunkSize + Overhead
)

func deriveKey(key, nonce []byte, info string, length int) (k []byte) {
	h := hkdf.New(sha256.New, key, nonce, []byte(info))
	k = make([]byte, length)
	if _, err := io.ReadFull(h, k); err != nil {
		panic("crypto: internal error: failed to read from HKDF: " + err.Error())
	}
	return
}

const encodedSizeLen = 4

func encodeLen(l int) []byte {
	b := make([]byte, encodedSizeLen)
	binary.LittleEndian.PutUint32(b, uint32(l))
	return b
}

func extract(p image.Image, i, j, size int) (int, int, []byte) {
	b := make([]byte, size)
	rect := p.Bounds()
	c := 0
	for ; i < rect.Max.Y; i++ {
		for ; j < rect.Max.X; j++ {
			if c == size {
				return i, j, b
			}
			pix := color.NRGBAModel.Convert(p.At(j, i)).(color.NRGBA)
			z := (pix.R << 6)
			z |= (pix.G << 6) >> 2
			z |= (pix.B << 6) >> 4
			z |= (pix.A << 6) >> 6
			b[c] = z
			c++
		}
		j = 0
	}
	return i, j, b
}

func bytesToLen(b []byte) int {
	return int(binary.LittleEndian.Uint32(b))
}

func Size(s int) int {
	return s + int(math.Ceil(float64(s)/ChunkSize)*Overhead)
}

func unseal(p image.Image, key []byte) (b []byte) {
	if key != nil {
		i, j, nonce := extract(p, 0, 0, NonceSize)
		k := deriveKey(key, nonce, infoEncryption, KeySize)
		s := NewAEADStream(k).Decrypt()
		i, j, firstBlock := extract(p, i, j, FullSize)
		dec := make([]byte, ChunkSize)
		s.XORKeyStream(dec, firstBlock)
		l := bytesToLen(dec[:encodedSizeLen]) + encodedSizeLen - ChunkSize
		_, _, data := extract(p, i, j, Size(l))
		b = make([]byte, l)
		s.XORKeyStream(b, data)
		return append(dec[encodedSizeLen:], b...)
	}
	i, j, lbuf := extract(p, 0, 0, encodedSizeLen)
	_, _, b = extract(p, i, j, bytesToLen(lbuf))
	return
}

const (
	infoEncryption  = "ChaCha20-Poly1305 encryption key"
	infoFiller      = "ChaCha20-Poly1305 randomness filler"
	infoFillerNonce = "ChaCha20-Poly1305 filler nonce"
	infoFillerStart = "RandomnessFiller"
)

func filler(key, nonce []byte, l int) (r []byte) {
	k := deriveKey(key, nonce, infoFiller, KeySize)
	n := deriveKey(key, nonce, infoFillerNonce, NonceSize)
	c, err := chacha20.NewUnauthenticatedCipher(k, n)
	if err != nil {
		panic(err)
	}
	var a, b [ChunkSize]byte
	copy(a[:], []byte(infoFillerStart))
	for len(r) < l {
		c.XORKeyStream(b[:], a[:])
		r = append(r, b[:]...)
		a = b
	}
	return
}

type AEADStream struct {
	c       cipher.AEAD
	nonce   [NonceSize]byte
	counter uint64
}

func NewAEADStream(key []byte) *AEADStream {
	c, err := chacha20poly1305.New(key)
	if err != nil {
		panic(err)
	}
	return &AEADStream{c: c}
}
func (aead *AEADStream) Encrypt() cipher.Stream {
	return &Encrypt{aead}
}
func (aead *AEADStream) Decrypt() cipher.Stream {
	return &Decrypt{aead}
}
func (aead *AEADStream) NonceSize() int {
	return aead.c.NonceSize()
}
func (aead *AEADStream) Overhead() int {
	return aead.c.Overhead()
}
func (aead *AEADStream) Seal(dst, _, plaintext, additionalData []byte) []byte {
	aead.inc()
	b := aead.c.Seal(dst, aead.nonce[:], plaintext, additionalData)
	return b
}
func (aead *AEADStream) Open(dst, _, ciphertext, additionalData []byte) ([]byte, error) {
	aead.inc()
	b, err := aead.c.Open(dst, aead.nonce[:], ciphertext, additionalData)
	return b, err
}
func (aead *AEADStream) inc() {
	put(&aead.nonce, aead.counter)
	aead.counter++
}

func put(nonce *[NonceSize]byte, n uint64) {
	binary.BigEndian.PutUint64((*nonce)[3:11], n)
}

type Encrypt struct {
	cipher.AEAD
}

func h(b []byte, v ...any) {
	fmt.Print(hex.EncodeToString(b) + " ")
	fmt.Println(v...)
}

func (e *Encrypt) XORKeyStream(dst, src []byte) {
	size := ChunkSize
	var buffer []byte
	var a [FullSize]byte
	for len(src) > 0 {
		if len(src) < size {
			size = len(src)
		}
		b := src[:size]
		src = src[size:]
		buf := e.Seal(a[:0], nil, b, nil)
		buffer = append(buffer, buf...)
	}
	copy(dst, buffer)
}

type Decrypt struct {
	cipher.AEAD
}

func (d *Decrypt) XORKeyStream(dst, src []byte) {
	size := FullSize
	var buffer []byte
	var a [ChunkSize]byte
	for len(src) > 0 {
		if len(src) < size {
			size = len(src)
		}
		b := src[:size]
		src = src[size:]
		buf, err := d.Open(a[:0], nil, b, nil)
		if err != nil {
			panic(err)
		}
		buffer = append(buffer, buf...)
	}
	copy(dst, buffer)
}

func encrypt(data, key []byte, max int) (enc []byte, err error) {
	nonce := make([]byte, NonceSize)
	_, err = io.ReadFull(rand.Reader, nonce)
	if err != nil {
		return
	}
	k := deriveKey(key, nonce, infoEncryption, KeySize)
	s := NewAEADStream(k).Encrypt()
	enc = make([]byte, Size(len(data)))
	s.XORKeyStream(enc, data)
	enc = append(append(nonce, enc...), filler(key, nonce, max)...)[:max]
	return
}

func seal(p image.Image, s, key []byte) (img image.Image, err error) {
	rect := p.Bounds()
	n := image.NewNRGBA(rect)
	c := 0
	arr := append(encodeLen(len(s)), s...)
	l := len(arr)
	max := (rect.Max.Y * rect.Max.X)
	if l > max {
		return nil, errors.New("not enough space in image")
	}
	if key != nil {
		arr, err = encrypt(arr, key, max)
		if err != nil {
			return
		}
		l = max
	}
	for i := 0; i < rect.Max.Y; i++ {
		for j := 0; j < rect.Max.X; j++ {
			pix := color.NRGBAModel.Convert(p.At(j, i)).(color.NRGBA)
			if c < l {
				z := arr[c]
				r := (pix.R &^ 3) | (z >> 6)
				g := (pix.G &^ 3) | (z<<2)>>6
				b := (pix.B &^ 3) | (z<<4)>>6
				a := (pix.A &^ 3) | (z<<6)>>6
				n.Set(j, i, color.NRGBA{r, g, b, a})
				c++
			} else {
				n.Set(j, i, pix)
			}
		}
	}
	return n, nil
}

var (
	inFile  string
	outFile string
	keyFile string
	secret  string
)

func readKey() (b []byte, err error) {
	if keyFile != "" {
		if keyFile == "-" {
			b, err = io.ReadAll(os.Stdin)
		} else {
			b, err = os.ReadFile(keyFile)
		}
	} else {
		if k := os.Getenv(EnvKey); k != "" {
			b = []byte(k)
		} else {
			return nil, nil
		}
	}
	if err != nil {
		return
	}
	return hex.DecodeString(string(b))
}

func errf(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

const EnvKey = "STEGO_KEY"

func main() {
	flag.StringVar(&inFile, "in", "-", "input file (default: stdin)")
	flag.StringVar(&keyFile, "key", "", "key file; used for encryption")
	flag.StringVar(&outFile, "out", "-", "output file, always png (default: stdout)")
	flag.StringVar(&secret, "secret", "", "secret file to hide")
	flag.Parse()
	if inFile == "-" && keyFile == "-" {
		errf(errors.New("can not read input and key files from stdin"))
	}
	key, err := readKey()
	if err != nil {
		errf(err)
	}
	var (
		in  = os.Stdin
		out = os.Stdout
	)
	if inFile != "-" {
		in, err = os.Open(inFile)
		if err != nil {
			errf(err)
		}
	}
	if outFile != "-" {
		out, err = os.OpenFile(outFile, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			errf(err)
		}
	}
	p, _, err := image.Decode(in)
	if err != nil {
		errf(err)
	}
	if secret != "" {
		b, err := os.ReadFile(secret)
		if err != nil {
			errf(err)
		}
		n, err := seal(p, b, key)
		if err != nil {
			errf(err)
		}
		err = png.Encode(out, n)
		if err != nil {
			errf(err)
		}
		return
	}
	_, err = out.Write(unseal(p, key))
	if err != nil {
		errf(err)
	}
}