main.go

package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"os"

	"golang.org/x/crypto/pbkdf2"
)

// ---------- Matrix Cipher ----------

type MatrixCipher struct {
	state [4][4]byte
	A     [4][4]byte
}

func NewMatrixCipher(key []byte) *MatrixCipher {
	var state [4][4]byte
	for i := 0; i < 4; i++ {
		for j := 0; j < 4; j++ {
			state[i][j] = key[i*4+j]
		}
	}
	// Fixed 4x4 matrix A (invertible modulo 256, used for diffusion)
	A := [4][4]byte{
		{2, 3, 1, 1},
		{1, 2, 3, 1},
		{1, 1, 2, 3},
		{3, 1, 1, 2},
	}
	return &MatrixCipher{state: state, A: A}
}

// XORStream encrypts/decrypts src and writes result to dst.
// Both dst and src must have the same length.
func (c *MatrixCipher) XORStream(dst, src []byte) {
	if len(dst) < len(src) {
		panic("dst too short")
	}
	blockSize := 16
	for i := 0; i < len(src); i += blockSize {
		n := blockSize
		if i+n > len(src) {
			n = len(src) - i
		}
		// Generate keystream from current state (row-major)
		ks := make([]byte, n)
		idx := 0
		for row := 0; row < 4 && idx < n; row++ {
			for col := 0; col < 4 && idx < n; col++ {
				ks[idx] = c.state[row][col]
				idx++
			}
		}
		// XOR
		for j := 0; j < n; j++ {
			dst[i+j] = src[i+j] ^ ks[j]
		}
		// Update state: state = state * A (mod 256)
		var newState [4][4]byte
		for row := 0; row < 4; row++ {
			for col := 0; col < 4; col++ {
				var sum uint16
				for k := 0; k < 4; k++ {
					sum += uint16(c.state[row][k]) * uint16(c.A[k][col])
				}
				newState[row][col] = byte(sum % 256)
			}
		}
		c.state = newState
	}
}

// ---------- MP4 parsing ----------

type MdatInfo struct {
	BoxOffset     int64 // offset of the 'mdat' box start
	HeaderSize    int64 // 8 or 16 bytes
	PayloadOffset int64 // offset of the payload (data inside the box)
	PayloadSize   int64 // size of the payload
	TotalBoxSize  int64 // size of the whole box (header + payload)
}

// findMdat scans the file for the 'mdat' box and returns its details.
func findMdat(data []byte) (*MdatInfo, error) {
	offset := 0
	for offset < len(data) {
		if offset+8 > len(data) {
			return nil, errors.New("truncated file: cannot read box header")
		}
		size := binary.BigEndian.Uint32(data[offset : offset+4])
		boxType := string(data[offset+4 : offset+8])

		var boxSize uint64
		var headerSize int
		if size == 1 {
			if offset+16 > len(data) {
				return nil, errors.New("truncated file: extended size missing")
			}
			boxSize = binary.BigEndian.Uint64(data[offset+8 : offset+16])
			headerSize = 16
		} else if size == 0 {
			// Box extends to end of file
			boxSize = uint64(len(data) - offset)
			headerSize = 8
		} else {
			boxSize = uint64(size)
			headerSize = 8
		}

		if boxType == "mdat" {
			return &MdatInfo{
				BoxOffset:     int64(offset),
				HeaderSize:    int64(headerSize),
				PayloadOffset: int64(offset + headerSize),
				PayloadSize:   int64(boxSize - uint64(headerSize)),
				TotalBoxSize:  int64(boxSize),
			}, nil
		}

		if size == 0 {
			// size==0 means this is the last box; break if we didn't find mdat
			break
		}
		offset += int(boxSize)
	}
	return nil, errors.New("mdat box not found")
}

// makeBoxHeader returns the binary header for an MP4 box.
func makeBoxHeader(size uint64, boxType string) []byte {
	buf := make([]byte, 16)
	if size < 0xFFFFFFFF {
		binary.BigEndian.PutUint32(buf[0:4], uint32(size))
		copy(buf[4:8], boxType)
		return buf[:8]
	}
	// Extended size
	binary.BigEndian.PutUint32(buf[0:4], 0xFFFFFFFF)
	copy(buf[4:8], boxType)
	binary.BigEndian.PutUint64(buf[8:16], size)
	return buf[:16]
}

// ---------- Embed & Extract ----------

// embedPDF reads an MP4 and a PDF, encrypts the PDF, and embeds it into the MP4's mdat box.
func embedPDF(mp4Path, pdfPath, outputPath, password string) error {
	mp4Data, err := os.ReadFile(mp4Path)
	if err != nil {
		return fmt.Errorf("read MP4: %w", err)
	}
	pdfData, err := os.ReadFile(pdfPath)
	if err != nil {
		return fmt.Errorf("read PDF: %w", err)
	}

	info, err := findMdat(mp4Data)
	if err != nil {
		return fmt.Errorf("find mdat: %w", err)
	}

	// Derive encryption key from password
	salt := []byte("MatrixEncryptionSalt")
	key := pbkdf2.Key([]byte(password), salt, 4096, 32, sha256.New)

	// Encrypt the PDF
	cipher := NewMatrixCipher(key)
	encryptedPDF := make([]byte, len(pdfData))
	cipher.XORStream(encryptedPDF, pdfData)

	// Prepare the data to append inside mdat
	magic := []byte("MPDP")
	origPayloadSize := make([]byte, 8)
	binary.BigEndian.PutUint64(origPayloadSize, uint64(info.PayloadSize))
	pdfSize := make([]byte, 8)
	binary.BigEndian.PutUint64(pdfSize, uint64(len(pdfData)))

	// Original mdat payload (video/audio data)
	originalPayload := mp4Data[info.PayloadOffset : info.PayloadOffset+info.PayloadSize]

	// New payload: original + magic + origSize + pdfSize + encryptedPDF
	newPayload := make([]byte, 0, len(originalPayload)+4+8+8+len(encryptedPDF))
	newPayload = append(newPayload, originalPayload...)
	newPayload = append(newPayload, magic...)
	newPayload = append(newPayload, origPayloadSize...)
	newPayload = append(newPayload, pdfSize...)
	newPayload = append(newPayload, encryptedPDF...)

	newTotalBoxSize := info.HeaderSize + int64(len(newPayload))
	newBoxHeader := makeBoxHeader(uint64(newTotalBoxSize), "mdat")

	// Reassemble the whole MP4 file
	result := make([]byte, 0, info.BoxOffset+int64(len(newBoxHeader))+int64(len(newPayload))+int64(len(mp4Data)-int(info.PayloadOffset+info.PayloadSize)))
	result = append(result, mp4Data[:info.BoxOffset]...)
	result = append(result, newBoxHeader...)
	result = append(result, newPayload...)
	// Append any trailing bytes that were after the original mdat box
	result = append(result, mp4Data[info.PayloadOffset+info.PayloadSize:]...)

	if err := os.WriteFile(outputPath, result, 0644); err != nil {
		return fmt.Errorf("write output: %w", err)
	}
	return nil
}

// extractPDF locates the hidden encrypted PDF in the MP4, decrypts it, and writes it back.
func extractPDF(mp4Path, pdfOutputPath, password string) error {
	mp4Data, err := os.ReadFile(mp4Path)
	if err != nil {
		return fmt.Errorf("read MP4: %w", err)
	}

	info, err := findMdat(mp4Data)
	if err != nil {
		return fmt.Errorf("find mdat: %w", err)
	}

	payload := mp4Data[info.PayloadOffset : info.PayloadOffset+info.PayloadSize]

	// Locate the magic marker inside the payload
	magic := []byte("MPDP")
	idx := bytes.LastIndex(payload, magic)
	if idx == -1 {
		return errors.New("hidden PDF not found (magic marker missing)")
	}
	if idx+4+8+8 > len(payload) {
		return errors.New("invalid hidden data format")
	}

	origPayloadSize := binary.BigEndian.Uint64(payload[idx+4 : idx+12])
	encryptedPDFSize := binary.BigEndian.Uint64(payload[idx+12 : idx+20])

	// Extract the encrypted PDF
	start := idx + 20
	end := start + int(encryptedPDFSize)
	if end > len(payload) {
		return errors.New("encrypted PDF data truncated")
	}
	encryptedPDF := payload[start:end]

	// Derive key
	salt := []byte("MatrixEncryptionSalt")
	key := pbkdf2.Key([]byte(password), salt, 4096, 32, sha256.New)

	// Decrypt (XOR is symmetric)
	cipher := NewMatrixCipher(key)
	decryptedPDF := make([]byte, len(encryptedPDF))
	cipher.XORStream(decryptedPDF, encryptedPDF)

	if err := os.WriteFile(pdfOutputPath, decryptedPDF, 0644); err != nil {
		return fmt.Errorf("write PDF: %w", err)
	}

	fmt.Printf("Original mdat payload size was %d bytes\n", origPayloadSize)
	return nil
}

// ---------- CLI ----------

func main() {
	mode := flag.String("mode", "", "embed or extract")
	mp4 := flag.String("mp4", "", "path to MP4 file (carrier)")
	pdf := flag.String("pdf", "", "path to PDF file (only for embed)")
	output := flag.String("output", "", "output file path (embedded MP4 for embed, decrypted PDF for extract)")
	password := flag.String("password", "", "encryption password")
	flag.Parse()

	if *mode == "" || *mp4 == "" || *output == "" || *password == "" {
		fmt.Println("Usage:")
		fmt.Println("  Embed:  program -mode embed -mp4 <carrier.mp4> -pdf <secret.pdf> -output <out.mp4> -password <pass>")
		fmt.Println("  Extract: program -mode extract -mp4 <embedded.mp4> -output <recovered.pdf> -password <pass>")
		flag.PrintDefaults()
		os.Exit(1)
	}

	var err error
	switch *mode {
	case "embed":
		if *pdf == "" {
			fmt.Println("error: -pdf is required for embed mode")
			os.Exit(1)
		}
		err = embedPDF(*mp4, *pdf, *output, *password)
	case "extract":
		err = extractPDF(*mp4, *output, *password)
	default:
		err = fmt.Errorf("invalid mode: %s (use 'embed' or 'extract')", *mode)
	}

	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %v\n", err)
		os.Exit(1)
	}
	fmt.Println("done")
}