/*
 * Guillaume Cottenceau (gc@mandrakesoft.com)
 *
 * Copyright 2000 MandrakeSoft
 *
 * This software may be freely redistributed under the terms of the GNU
 * public license.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 */

/*
 * Portions from Erik Troan (ewt@redhat.com)
 *
 * Copyright 1996 Red Hat Software
 *
 */

#ifndef _GNU_SOURCE
#define _GNU_SOURCE
#endif
#include <stdlib.h>
#include <unistd.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <ctype.h>
#include <stdio.h>
#include <dirent.h>
#include <sys/types.h>
#include <sys/mount.h>
#include <sys/poll.h>
#include <sys/sysinfo.h>
#include <sys/mman.h> /* memfd_create */

#include "stage1.h"
#include "log.h"
#include "mount.h"
#include "frontend.h"
#include "automatic.h"
#include <errno.h>
#include <linux/loop.h>
#include "lomount.h"
#include "tools.h"
#include "modules.h"
#include "sha256.h"
#ifdef SPAWN_SPLASH
#include "init.h"
#endif
#include "udev.h"

/* If we have more than that amount of memory (in bytes), we assume we can load the second stage as a ramdisk */
#define MEM_LIMIT_RAMDISK (256 * 1024 * 1024)

/* If we have more than that amount of memory (in bytes), we assume we can load the rescue as a ramdisk */
#define MEM_LIMIT_RESCUE (128 * 1024 * 1024)

static struct param_elem params[50];
static int param_number = 0;
static unsigned long cmdline_ramdisk_size = 0;

void process_cmdline(void)
{
	char buf[512];
	int fd, size, i;

	log_message("opening /proc/cmdline... ");

	if ((fd = open("/proc/cmdline", O_RDONLY)) == -1)
		fatal_error("could not open /proc/cmdline");

	size = read(fd, buf, sizeof(buf));
	buf[size-1] = '\0'; // -1 to eat the \n
	close(fd);

	log_message("\t%s", buf);

	i = 0;
	while (buf[i] != '\0') {
		char *name, *value = NULL;
		int j = i;
		while (buf[i] != ' ' && buf[i] != '=' && buf[i] != '\0')
			i++;
		if (i == j) {
			i++;
			continue;
		}
		name = memdup(&buf[j], i-j + 1);
		name[i-j] = '\0';

		if (buf[i] == '=') {
			int k = i+1;
			i++;
			while (buf[i] != ' ' && buf[i] != '\0')
				i++;
			value = memdup(&buf[k], i-k + 1);
			value[i-k] = '\0';
		}

		params[param_number].name = name;
		params[param_number].value = value;
		param_number++;
		if (!strcmp(name, "expert")) set_param(MODE_EXPERT);
		else if (!strcmp(name, "changedisk")) set_param(MODE_CHANGEDISK);
		else if (!strcmp(name, "updatemodules")) set_param(MODE_UPDATEMODULES);
		else if (!strcmp(name, "rescue")) set_param(MODE_RESCUE);
		else if (!strcmp(name, "splash")) set_param(MODE_SPLASH);
		else if (!strcmp(name, "live")) set_param(MODE_LIVE);
		else if (!strcmp(name, "stagename")) set_param(MODE_STAGENAME);
		else if (!strcmp(name, "lowmem")) set_param(MODE_LOWMEM);
		else if (!strcmp(name, "hash")) set_param(MODE_VERIFICATION);
		else if (!strcmp(name, "automatic")) {
			set_param(MODE_AUTOMATIC);
			grab_automatic_params(value);
		} else if (!strcmp(name, "ramdisk_size")) {
			long tmp;
			char *endptr = NULL;
			errno = 0;
			tmp = strtol(value, &endptr, 0);
			if (errno == 0 && endptr && *endptr == '\0' && tmp > 0) {
				/* ramdisk_size in cmdline given in Kbs, convert
				 * its value to bytes */
				cmdline_ramdisk_size = (unsigned long)tmp * 1024;
			}
		}
		if (buf[i] == '\0')
			break;
		i++;
	}

	log_message("\tgot %d args", param_number);
}


int stage1_mode = 0;

int get_param(int i)
{
#ifdef SPAWN_INTERACTIVE
	static int fd = 0;
	char buf[5000];
	char * ptr;
	int nb;

	if (fd <= 0) {
		fd = open(interactive_fifo, O_RDONLY);
		if (fd == -1)
			return (stage1_mode & i);
		fcntl(fd, F_SETFL, O_NONBLOCK);
	}

	if (fd > 0) {
		if ((nb = read(fd, buf, sizeof(buf))) > 0) {
			buf[nb] = '\0';
			ptr = buf;
			while ((ptr = strstr(ptr, "+ "))) {
				if (!strncmp(ptr+2, "expert", 6)) set_param(MODE_EXPERT);
				if (!strncmp(ptr+2, "rescue", 6)) set_param(MODE_RESCUE);
				ptr++;
			}
			ptr = buf;
			while ((ptr = strstr(ptr, "- "))) {
				if (!strncmp(ptr+2, "expert", 6)) unset_param(MODE_EXPERT);
				if (!strncmp(ptr+2, "rescue", 6)) unset_param(MODE_RESCUE);
				ptr++;
			}
		}
	}
#endif

	return (stage1_mode & i);
}

char * get_param_valued(char *param_name)
{
	int i;
	for (i = 0; i < param_number ; i++)
		if (!strcmp(params[i].name, param_name))
			return params[i].value;

	return NULL;
}

void set_param_valued(char *param_name, char *param_value)
{
	params[param_number].name = param_name;
	params[param_number].value = param_value;
	param_number++;
}

void set_param(int i)
{
	stage1_mode |= i;
	if (i == MODE_RESCUE) {
		set_param_valued("stagename", "rescue");
		set_param(MODE_STAGENAME);
	}
}

void unset_param(int i)
{
	stage1_mode &= ~i;
}

// warning, many things rely on the fact that:
// - when failing it returns 0
// - it stops on first non-digit char
int charstar_to_int(char * s)
{
	int number = 0;
	while (*s && isdigit(*s)) {
		number = (number * 10) + (*s - '0');
		s++;
	}
	return number;
}

/* Returns total memory in bytes */
unsigned long total_memory(void)
{
	struct sysinfo info;
	unsigned long value;

    if (sysinfo(&info)) {
		log_message( "sysinfo: Can't get total memory: %s", strerror(errno));
        return 0;
    }

	value = info.totalram * info.mem_unit;

	return value;
}

unsigned long get_ramdisk_size(const char *ramdisk_path)
{
	/* If path is given then use it and return 0 if stat failed */
	if (ramdisk_path) {
		struct stat statbuf;

		return stat(ramdisk_path, &statbuf) ? 0 : (unsigned long)statbuf.st_size;
	}

	/* Use ramdisk_size from cmdline if given */
	if (cmdline_ramdisk_size)
		return cmdline_ramdisk_size;

	/* Fallback to hardcoded values as last resort */
	return IS_RESCUE ? MEM_LIMIT_RESCUE : MEM_LIMIT_RAMDISK;
}


int ramdisk_possible(unsigned long ramdisk_size)
{
	unsigned long tm = total_memory();

	log_message("Total Memory: %u Mbytes", BYTES2MB(tm));

	/* Assume we need at least twice of ramdisk size */
	if (total_memory() >= ramdisk_size * 2)
		return 1;
	else {
		log_message("warning, ramdisk (%u Mb) is not possible due to low mem!", BYTES2MB(ramdisk_size));
		return 0;
	}
}

static void save_stuff_for_rescue(void)
{
	char * file = "/etc/resolv.conf";
	char buf[5000];
	int fd_r, fd_w, i;
	char location[100];

	if ((fd_r = open(file, O_RDONLY)) < 0) {
		log_message("can't open %s for read", file);
		return;
	}
	strcpy(location, STAGE2_LOCATION);
	strcat(location, file);
	if ((fd_w = open(location, O_WRONLY)) < 0) {
		log_message("can't open %s for write", location);
		close(fd_r);
		return;
	}
	if ((i = read(fd_r, buf, sizeof(buf))) <= 0) {
		log_message("can't read from %s", file);
		close(fd_r); close(fd_w);
		return;
	}
	if (write(fd_w, buf, i) != i)
		log_message("can't write %d bytes to %s", i, location);
	close(fd_r); close(fd_w);
	log_message("saved file %s for rescue (%d bytes)", file, i);
}

/**
 * Loop mount ISO loaded to RAM, and open stage2
 *
 * @param ram_fd file descriptor of ISO image
 * @return file descriptor of stage2 or -1 on error
 */
static int open_stage2_from_iso(int iso_fd)
{
	int stage2_fd = -1;
	const char *stage2 = get_ramdisk_path(IMAGE_LOCATION); /* /image/live */
	if (!stage2) {
		log_message("%s: get_ramdisk_size: got NULL", __func__);
		goto out;
	}
	if (do_losetup_fd(LOMOUNT_DEVICE, iso_fd, NULL) < 0) {
		log_message("%s: could not setup loopback for iso_fd", __func__);
		goto out_free;
	}
	if (my_mount(LOMOUNT_DEVICE, IMAGE_LOCATION, "iso9660", 0)) {
		log_message("%s: failed to mount ISO loopback", __func__);
		goto out_del_loop;
	}
	stage2_fd = open(stage2, O_RDONLY);
	if (stage2_fd < 0) {
		log_message("%s: failed to open %s", __func__, stage2);
		goto out_umount;
	}
	free((char *)stage2);
	return stage2_fd;

out_umount:
	umount(IMAGE_LOCATION);
out_del_loop:
	del_loop(LOMOUNT_DEVICE);
out_free:
	free((char *)stage2);
out:
	return -1;
}

/**
 * @brief copy exactly len bytes from buffer to file descriptor
 */
static int write_exactly(int fd, const char *buf, size_t len)
{
	ssize_t dl;
	while (len > 0) {
		dl = write(fd, buf, len);
		if (dl < 0) {
			log_message("%s: write: %s", __func__, strerror(errno));
			return -1;
		}
		buf += (size_t)dl;
		len -= (size_t)dl;
	}
	return 0;
}

/*
 * @brief Copy size bytes from src_fd to dst_fd
 * @param exact if non-zero check if exactly size bytes have been copied
 * @return 0 on success, -1 on error
 * @note if size == 0 keep copying until EOF on src_fd
 */
static int copy_loop(int dst_fd, int src_fd, unsigned long size)
{
	char buf[32768];
	unsigned long bytes_written = 0;
	ssize_t dl;
	for (;;) {
		dl = read(src_fd, buf, sizeof(buf));
		if (dl < 0) {
			log_message("%s: read: %s", __func__, strerror(errno));
			return -1;
		} else if (dl == 0) {
			break;
		}
		if (write_exactly(dst_fd, buf, (size_t)dl) < 0) {
			return -1;
		}
		bytes_written += (size_t)dl;
		update_progression((int)BYTES2MB(bytes_written));
	}
	if (size > 0 && bytes_written != size) {
		log_message("%s: expected %lu bytes, copied %lu instead\n",
			__func__, size, bytes_written);
		errno = EIO;
		return -1;
	}
	return 0;
}

int make_ramfd(unsigned long size)
{
	int ramfd;
	ramfd = memfd_create("image", MFD_ALLOW_SEALING);
	if (ramfd < 0) {
		log_perror("memfd_create");
		return -1;
	}

	/* By default memfd has no size limit. Set it to avoid OOM */
	if (!size) {
		/* XXX: when downloading an image via FTP size can
		 * be 0/unknown (the protocol makes it extremely
		 * difficult to figure out the file size).
		 * Set the limit to of 1/2 RAM in this case.
		 */
		size = total_memory()/2;
	}
	if (ftruncate(ramfd, (off_t)size) < 0) {
		log_perror("ftruncate ramfd");
		goto out;
	}
	if (fcntl(ramfd, F_ADD_SEALS, F_SEAL_GROW) < 0) {
		log_perror("fcntl ramfd F_SEAL_GROW");
		goto out;
	}
	return ramfd;
out:
	close(ramfd);
	return -1;
}

enum return_type load_ramdisk_or_iso(int source_fd, unsigned long size, int is_iso)
{
	int ram_fd;
	char * wait_msg = "Loading program into memory...";

	ram_fd = make_ramfd(size);
	if (ram_fd < 0) {
		stg1_error_message("Couldn't allocate memory for image.");
		return RETURN_ERROR;
	}

	init_progression(wait_msg, (int)BYTES2MB(size));
	if (copy_loop(ram_fd, source_fd, size) < 0) {
		log_message("%s: failed to load image to RAM", __func__);
		close(ram_fd);
		remove_wait_message();
		stg1_error_message("Could not load second stage ramdisk. "
				   "This is probably an hardware error while reading the data. "
				   "(this may be caused by a hardware failure or a Linux kernel bug)");
		return RETURN_ERROR;
	}
	end_progression();

	if (is_iso) {
		int stage2_fd = open_stage2_from_iso(ram_fd);
		if (stage2_fd < 0) {
			close(ram_fd);
			stg1_error_message("Could not open stage2 from ISO in RAM");
			return RETURN_ERROR;
		}
		/* got a (mounted) loopback device backed by ram_fd,
		 * ram_fd itself is not necessary any more */
		close(ram_fd);
		/* to setup loopback device from stage2_fd */
		ram_fd = stage2_fd;
	}

	if (do_losetup_fd(LIVE_DEVICE, ram_fd, NULL) < 0) {
		close(ram_fd);
		if (is_iso) {
			umount(IMAGE_LOCATION);
			del_loop(LOMOUNT_DEVICE);
		}
		stg1_error_message("Could not setup loopback for 2nd stage");
		return RETURN_ERROR;
	}

	/* got a loopback device, don't need ram_fd any more */
	close(ram_fd);

	if (my_mount(LIVE_DEVICE, STAGE2_LOCATION, STAGE2FS, 0)) {
		stg1_error_message("Failed to mount stage2");
		if (is_iso) {
			umount(IMAGE_LOCATION);
			del_loop(LOMOUNT_DEVICE);
		}
		return RETURN_ERROR;
	}

	set_param(MODE_RAMDISK);

	if (IS_RESCUE) {
		save_stuff_for_rescue();
		if (umount(STAGE2_LOCATION)) {
			log_message("rescue: failed to umount " STAGE2_LOCATION);
			return RETURN_ERROR;
		}
		if (is_iso) {
			/* XXX: should I do anything special here? */
		}
		return RETURN_OK; /* fucksike, I lost several hours wondering why the kernel won't see the rescue if it is alreay mounted */
	}

	return RETURN_OK;
}

enum return_type load_ramdisk_fd(int source_fd, unsigned long size) {
	return load_ramdisk_or_iso(source_fd, size, 0);
}

int splash_verbose()
{
#ifdef SPAWN_SPLASH
	if (!IS_SPLASH) return(1);

	struct stat pst;
	if (stat("/bin/plymouth", &pst)) return(1);

	char * av[] = 	{ "/bin/plymouth" , "plymouth" , "quit", NULL };
	log_message( "%s: %s\n", av[0], av[2] );
	return 	spawn(av);
#endif
}

int update_splash( char * state )
{
#ifdef SPAWN_SPLASH
	if (!IS_SPLASH) return(1);

	struct stat pst;
	if (stat("/bin/plymouth", &pst)) return(1);

	char * av[] = 	{ "/bin/plymouth" , "plymouth" , "--update" , state, NULL };
	log_message( "%s: %s\n", av[0], av[3] );
	return 	spawn(av);
#endif
}

char * get_ramdisk_realname(void)
{
	char img_name[500];
	char * stg2_name = get_param_valued("stagename");
	char * begin_img = RAMDISK_LOCATION;

	if (!stg2_name)
		stg2_name = "altinst";

	strcpy(img_name, begin_img);
	strcat(img_name, stg2_name);

	return strdup(img_name);
}

char * get_ramdisk_path(const char *mount_path)
{
    char img_name[500];
    char * st2_path = get_ramdisk_realname();

	/* FIXME */
	strcpy(img_name, mount_path ? mount_path : IMAGE_LOCATION);
    strcat(img_name, st2_path);
	free(st2_path);
    return strdup(img_name);
}

char * get_uid_file_path(const char *mount_path, const char *uid_file)
{
	char file_name[500];

	/* return NULL for empty file uid */
	if (!uid_file || uid_file[0] == '\0')
		return NULL;
	/* FIXME */
	strcpy(file_name, mount_path ? mount_path : IMAGE_LOCATION "/");
	strcat(file_name, uid_file);
	return strdup(file_name);
}

/* This function is used to protect against stage2 (aka ramdisk) spoofing */
enum return_type verify_ramdisk_digest(const char *filename, const char *sha256_hash)
{
	if (!sha256_hash)
		return RETURN_ERROR;

	int fd;
	struct stat st;

	if ((fd = open(filename, O_RDONLY)) < 0 || fstat(fd, &st)) {
		log_message("verify_ramdisk_digest: %s: %m", filename);
		close(fd);
		return RETURN_ERROR;
	}

	init_progression("Verifying stage2 authenticity...", (int)BYTES2MB(st.st_size));

	sha256_context ctx;
	sha256_starts(&ctx);

	ssize_t bytes_read = 0, total_bytes_read = 0;
	unsigned char buffer[8192];

	while ((bytes_read = read(fd, buffer, sizeof(buffer)))) {
		if (bytes_read < 0) {
			if (EINTR == errno)
				continue;
			log_message("verify_ramdisk_digest: %s: %m", filename);
			close(fd);
			remove_wait_message();
			return RETURN_ERROR;
		}
		sha256_update(&ctx, buffer, bytes_read);
		total_bytes_read += bytes_read;
		update_progression((int)BYTES2MB(total_bytes_read));
	}

	close(fd);
	end_progression();

	uint8 digest[256/8];
	sha256_finish(&ctx, digest);

	const char *hex = "0123456789abcdef";
	unsigned i;
	char computed_hash[sizeof(digest)*2 + 1];
	char *dest = computed_hash;

	for (i = 0; i < sizeof(digest); ++i) {
		*(dest++) = hex[(digest[i] >> 4) & 0xf];
		*(dest++) = hex[digest[i] & 0xf];
	}
	*dest = '\0';

	return strcmp(computed_hash, sha256_hash) ? RETURN_ERROR : RETURN_OK;
}

enum return_type load_ramdisk(char *ramdisk_path, unsigned long size)
{
	int st2_fd;
	char *img_name;
	enum return_type rc = RETURN_ERROR;

	img_name = ramdisk_path ?: get_ramdisk_path(NULL);

	log_message("trying to load %s as a ramdisk", img_name);

	st2_fd = open(img_name, O_RDONLY); /* to be able to see the progression */

	if (st2_fd == -1) {
		log_message("open ramdisk file (%s) failed", img_name);
		stg1_error_message("Could not open compressed ramdisk file (%s).", img_name);
		goto done;
	}

	if (size == 0) {
		size = get_ramdisk_size(img_name);
		if (size == 0)
			goto done;
	}
	rc = load_ramdisk_fd(st2_fd, size);

done:
	if (img_name != ramdisk_path)
		free(img_name);

	return rc;
}

/* pixel's */
void * memdup(void *src, size_t size)
{
	void * r;
	r = malloc(size);
	if(!r)
		perror("malloc");
	memcpy(r, src, size);
	return r;
}


static char ** my_env = NULL;
static int env_size = 0;

void handle_env(char ** env)
{
	char ** ptr = env;
	while (ptr && *ptr) {
		ptr++;
		env_size++;
	}
	my_env = malloc(sizeof(char *) * 100);
	if(!my_env)
		perror("malloc");
	memcpy(my_env, env, sizeof(char *) * (env_size+1));
}

char ** grab_env(void) {
	return my_env;
}

void add_to_env(char * name, char * value)
{
	char tmp[500];
	sprintf(tmp, "%s=%s", name, value);
	my_env[env_size] = strdup(tmp);
	env_size++;
	my_env[env_size] = NULL;
}

char * get_from_env(const char * key, const char * const * env)
{
	int i=0;

	if (key == NULL || env == NULL)
		return NULL;

	while(env[i])
	{
		if (strncmp(env[i], key, strlen(key)) == 0) {
			return strdup(env[i]+strlen(key)+1);
		}
		i++;
	}

	return NULL;
}

int pass_env(int fd)
{
	char ** ptr = my_env;
	char *s;
	int i = 0;

	while (ptr && *ptr) {
		s = *ptr;
		while(*s++);
		i = write(fd, *ptr, s - *ptr);
		ptr++;
	}
	return i;
}

char ** list_directory(char * direct)
{
	char * tmp[50000]; /* in /dev there can be many many files.. */
	int i = 0;
	struct dirent *ep;
	DIR *dp = opendir(direct);
	while (dp && (ep = readdir(dp))) {
		if (strcmp(ep->d_name, ".") && strcmp(ep->d_name, "..")) {
			tmp[i] = strdup(ep->d_name);
			i++;
		}
	}
	if (dp)
		closedir(dp);
	tmp[i] = NULL;
	return memdup(tmp, sizeof(char*) * (i+1));
}


int string_array_length(char ** a)
{
	int i = 0;
	if (!a)
		return -1;
	while (a && *a) {
		a++;
		i++;
	}
	return i;
}

/* setup loop device from file descriptor targfd */
int do_losetup_fd(char *device, int targfd, const char *name)
{
	int i, loopfd;
	struct loop_info loopInfo;
	memset(&loopInfo, 0, sizeof(loopInfo));

	my_insmod("loop", NULL);
	/* wait for udev's dust settles down */
	for (i=3; i && (loopfd = open(device, O_RDONLY)) < 0; --i, sleep(1));
	if ( loopfd < 0 )
        {
		log_message( "losetup: error opening %s: %s", device, strerror(errno) );
		return -1;
	}

	if ( ioctl(loopfd, LOOP_SET_FD, targfd) < 0 )
	{
		log_message( "losetup: error setting up loopback device: %s", strerror(errno) );
		close( loopfd );
		return -1;
	}
	/* Note: LOOP_SET_FD triggers change event. While processing it
	 * udev reads the loop device with builtin blkid. This can race
	 * with subsequent access by kernel due to LOOP_SET_STATUS (or
	 * mounting the loop device). Therefore give udev a chance to
	 * process the event without concurrent access to the loop device
	 */
	udev_settle();

	if (name) {
		strcpy(loopInfo.lo_name, name);

		if ( ioctl(loopfd, LOOP_SET_STATUS, &loopInfo) < 0 )
			log_message( "losetup: error setting up loopback device: %s", strerror(errno) );

		/* Same here: LOOP_SET_STATUS triggers change event, give udev
		 * a chance to process it without concurrent access to the loop
		 * device. In other words, prevent the caller of this function
		 * from mounting the device while udev still running blkid on it
		 */
		udev_settle();
	}

	close( loopfd );
	/* udev might be monitoring loopback device (with fanotify) and
	 * will synthesize a change event on close. Give udev some time
	 * to process without racing with the caller of this function
	 * (which is going to mount the newly configured loopback device)
	 */
	udev_settle();
	return 0;
}

int do_losetup(char * device, char * target)
{
	int ret = 0, targfd;
	targfd = open( target, O_RDONLY );
	if ( targfd < 0 )
	{
		log_message( "losetup: error opening %s: %s", target, strerror(errno) );
		return -1;
	}
	ret = do_losetup_fd(device, targfd, target);
	close(targfd);
	return ret;
}