main.go

package main

import (
	"flag"
	"fmt"
	"image"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"os"

	"steg-go/server"
	"steg-go/stego"
)

func main() {
	// Check if any arguments were provided
	if len(os.Args) == 1 {
		// No arguments - launch web server
		port := 0 // Use OS-assigned port (any available)
		s := server.NewServer(port, embedFS)
		if err := s.Run(); err != nil {
			fmt.Fprintf(os.Stderr, "Server error: %v\n", err)
			os.Exit(1)
		}
		return
	}

	// CLI mode
	runCLI()
}

func runCLI() {
	var inputImage string
	var output string
	var files []string

	flag.StringVar(&inputImage, "i", "", "Input image file")
	flag.StringVar(&output, "o", "", "Output PNG file (encode) or directory (extract)")
	flag.Func("f", "Files or directories to hide (can be used multiple times)", func(s string) error {
		files = append(files, s)
		return nil
	})

	flag.Parse()

	// Collect remaining arguments after flags as files
	files = append(files, flag.Args()...)

	if inputImage == "" || output == "" {
		fmt.Fprintf(os.Stderr, "Usage: %s -i <input_image> [-f <file1> -f <file2> ...] -o <output>\n", os.Args[0])
		fmt.Fprintf(os.Stderr, "\nEncode mode: -i <image> -f <files...> -o <output.png>\n")
		fmt.Fprintf(os.Stderr, "Extract mode: -i <image> -o <output_directory>\n")
		fmt.Fprintf(os.Stderr, "\nGUI mode: Run without any arguments\n")
		os.Exit(1)
	}

	// Load input image
	img, err := loadImage(inputImage)
	if err != nil {
		fmt.Fprintf(os.Stderr, "Error loading image: %v\n", err)
		os.Exit(1)
	}

	// Check if this is extract mode (no files specified and output is a directory)
	outputStat, err := os.Stat(output)
	isExtractMode := len(files) == 0 && (err == nil && outputStat.IsDir())

	if isExtractMode {
		// Extract mode
		fmt.Println("Extract mode: extracting hidden files from image...")
		data, err := stego.DecodeData(img)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error decoding data: %v\n", err)
			os.Exit(1)
		}

		if len(data) == 0 {
			fmt.Println("No hidden data found in image")
			os.Exit(0)
		}

		fmt.Printf("Decoded hidden payload: %s\n", stego.FormatBytes(uint64(len(data))))

		err = stego.ExtractTarGz(data, output)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error extracting archive: %v\n", err)
			os.Exit(1)
		}

		fmt.Printf("Successfully extracted files to %s\n", output)
	} else {
		// Encode mode
		if len(files) == 0 {
			fmt.Println("Warning: No files to hide, creating image with no hidden data")
		}

		// Create tar.gz archive of files to hide
		var archiveData []byte
		if len(files) > 0 {
			archiveData, err = stego.CreateTarGzWithStructure(files)
			if err != nil {
				fmt.Fprintf(os.Stderr, "Error creating archive: %v\n", err)
				os.Exit(1)
			}
			fmt.Printf("Created archive with preserved directory structure: %s\n", stego.FormatBytes(uint64(len(archiveData))))
		}

		// Encode data into image using LSB steganography
		encodedImg, err := stego.EncodeData(img, archiveData)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error encoding data: %v\n", err)
			os.Exit(1)
		}

		// Save output PNG
		err = savePNG(encodedImg, output)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error saving PNG: %v\n", err)
			os.Exit(1)
		}

		fmt.Printf("Successfully saved steganographic image to %s\n", output)
	}
}

// loadImage loads an image from a file
func loadImage(path string) (image.Image, error) {
	file, err := os.Open(path)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	img, _, err := image.Decode(file)
	if err != nil {
		return nil, err
	}

	return img, nil
}

// savePNG saves an image as PNG
func savePNG(img image.Image, path string) error {
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	return png.Encode(file, img)
}