#include <fcntl.h>
#include <linux/limits.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <elf_c_types.h>
#include "elfio_c_wrapper.h"
#include <unistd.h>
#include "elf_package.h"

#define LOG_OUT stdout
#define ELF_GET_PACK_SECTION(pelfio, sec, alloc)             \
	sec = elfio_get_section_by_name(pelfio, PACK_NAME);  \
	if (!sec && alloc) {                                 \
		fprintf(LOG_OUT, "will add pack section\n"); \
		sec = elfio_add_section(pelfio, PACK_NAME);  \
		elfio_section_set_type(sec, SHT_NOTE);       \
	}

static bool inline is_elf_pack(Elf_Word type)
{
	if (type >= NT_PACK_NONE && type <= NT_PACK_END)
		return true;
	return false;
}

static pelfio_t elf_open(char *filename)
{
	pelfio_t pelfio = elfio_new();
	bool ret = true;
	char msg[128];

	ret = elfio_load(pelfio, filename);
	if (!ret) {
		fprintf(LOG_OUT, "Can't load ELF file\n");
		return NULL;
	}

	ret = elfio_validate(pelfio, msg, 128);
	if (!ret) {
		fprintf(LOG_OUT, "ELF file is invalid\n");
		return NULL;
	}

	return pelfio;
}

static void elf_close(pelfio_t pelfio)
{
	elfio_delete(pelfio);
}

static pnote_t elf_get_pack_note(pelfio_t pelfio, bool alloc)
{
	psection_t sec = NULL;

	ELF_GET_PACK_SECTION(pelfio, sec, alloc);
	if (!sec) {
		fprintf(LOG_OUT, "pack section get fail\n");
		return NULL;
	}

	return elfio_note_section_accessor_new(pelfio, sec);
}

static bool elf_data_exist(pnote_t note, char *name)
{
	int noteno = elfio_note_get_notes_num(note);
	Elf_Word type;
	char gname[PACK_NAME_MAX_LEN];
	int ret = true;
	char *desc;
	Elf_Word descSize = 128;

	for (int i = 0; i < noteno; i++) {
		ret = elfio_note_get_note(note, i, &type, gname,
					  PACK_NAME_MAX_LEN, (void **)&desc,
					  &descSize);
		if (!ret || !is_elf_pack(type)) {
			continue;
		}
		if (strcmp(gname, name) == 0) {
			return true;
		}
	}
	return false;
}

static bool elf_add_data(char *filename, char *data, int64_t size, char *name,
			 Elf_Word type)
{
	pelfio_t pelfio = elf_open(filename);
	pnote_t note = NULL;
	bool ret = true;

	if (!pelfio)
		return false;

	note = elf_get_pack_note(pelfio, true);
	if (elf_data_exist(note, name)) {
		fprintf(LOG_OUT, "The data[%s] is exist\n", name);
	}
	elfio_note_add_note(note, type, name, data, size);

	elfio_save(pelfio, filename);

close:
	elf_close(pelfio);
	return ret;
}

static Elf_Word elf_get_pack_type(char *file)
{
	Elf_Word type = 0;
	pelfio_t pelfio = elf_open(file);

	if (!pelfio) {
		goto out;
	}
	type = elfio_get_type(pelfio);
	elf_close(pelfio);
out:
	return NT_PACK_TYPE(type);
}

int elf_pack(char *soufile, char *file)
{
	char *name = NULL;
	int fd = -1;
	ssize_t size = 0;
	char *data = NULL;
	int ret = 0;
	Elf_Word type = elf_get_pack_type(file);

	fd = open(file, O_RDWR);
	if (fd < 0) {
		fprintf(LOG_OUT, "Error opening file: %s\n", file);
		return -1;
	}
	size = lseek(fd, 0, SEEK_END);
	lseek(fd, 0, SEEK_SET);
	data = calloc(1, size);
	if (!data) {
		ret = -1;
		fprintf(LOG_OUT, "malloc fail for: %s\n", file);
		goto close;
	}
	if (read(fd, data, size) != size) {
		ret = -1;
		fprintf(LOG_OUT, "read fail: %s\n", file);
		goto close;
	}

	name = strrchr(file, '/');
	if (name)
		name = name + 1;
	else
		name = file;

	if (!elf_add_data(soufile, data, size, name, type)) {
		ret = -1;
		fprintf(LOG_OUT, "Add data fail: %s\n", file);
	}

close:
	if (data)
		free(data);
	close(fd);
	return ret;
}

static bool is_dir(char *dir)
{
	struct stat st;
	if (stat(dir, &st) == 0 && S_ISDIR(st.st_mode))
		return true;
	return false;
}

static bool create_file(char *file, char *data, int64_t size, int mode)
{
	int fd = -1;
	int ret = true;

	fd = open(file, O_RDWR | O_CREAT | O_TRUNC, mode);
	if (fd < 0) {
		fprintf(LOG_OUT, "Error opening file: %s\n", file);
		return false;
	}
	if (write(fd, data, size) != size) {
		fprintf(LOG_OUT, "write fail: %s\n", file);
		ret = false;
	}

	close(fd);
	return ret;
}

static Elf_Word elf_get_file_mode(Elf_Word type)
{
	switch (type) {
	case NT_PACK_EXEC:
		return ACCESSPERMS;
	default:
		return DEFFILEMODE;
	}
}

static bool elf_note_unpack(pnote_t note, char *dir)
{
	int noteno = elfio_note_get_notes_num(note);
	Elf_Word type;
	char gname[PACK_NAME_MAX_LEN];
	int ret = true;
	char *desc;
	char PATH[PATH_MAX];
	Elf_Word descSize = 0;

	for (int i = 0; i < noteno; i++) {
		ret = elfio_note_get_note(note, i, &type, gname,
					  PACK_NAME_MAX_LEN, (void **)&desc,
					  &descSize);
		if (!ret || !is_elf_pack(type)) {
			fprintf(LOG_OUT, "get note %d fail\n", i);
			continue;
		}
		sprintf(PATH, "%s/%s", dir, gname);
		if (!create_file(PATH, desc, descSize,
				 elf_get_file_mode(type))) {
			fprintf(LOG_OUT, "create file %s fail\n", PATH);
			continue;
		}
	}
	return ret;
}

int elf_unpack(char *soufile, char *dir)
{
	pelfio_t pelfio = NULL;
	pnote_t note = NULL;
	int ret = 0;

	if (!is_dir(dir)) {
		fprintf(LOG_OUT, "dir[%s] is not exist\n", dir);
		return -1;
	}
	pelfio = elf_open(soufile);
	if (!pelfio) {
		fprintf(LOG_OUT, "open elf file fail\n");
		return -1;
	}
	note = elf_get_pack_note(pelfio, false);
	if (!note) {
		fprintf(LOG_OUT, "get pack note fail\n");
		ret = -1;
		goto close;
	}
	ret = !elf_note_unpack(note, dir);

close:
	elf_close(pelfio);
	return ret;
}