#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <string.h>
#include <sys/stat.h>

#define FLAG "UMASS{f1l3f0rm4754r3h4rd_r3m3mb3rh34r7bl33d}"

int main(int argc, char **argv)
{
		struct stat st;
		const char *path_in, *path_out, *data;
		char *buf, *output, *flag;
		size_t file_size, data_len;
		FILE *fp;
		int i;

		int32_t data_start;
		uint32_t compression, image_size;

		flag = malloc(strlen(FLAG) + 1);
		strcpy(flag, FLAG);

		if (argc != 4) {
				fprintf(stderr, "usage: %s [in] [out] [data]\n", argv[0]);
				return 1;
		}

		path_in = argv[1]; path_out = argv[2]; data = argv[3];

		stat(path_in, &st);
		file_size = st.st_size;

		buf = malloc(file_size);

		fp = fopen(path_in, "rb");
		if (fp == NULL) {
				fprintf(stderr, "%s: no such file\n", path_in);
				return 2;
		}

		fread(buf, file_size, 1, fp);
		fclose(fp);

		// Here's where the actual parsing
		if (buf[0] != 'B' || buf[1] != 'M') {
				fprintf(stderr, "%s: not a valid BMP\n", path_in);
				return 3;
		}

		compression = *((uint32_t *) (buf + 0x1e));
		if (compression != 0) {
				fprintf(stderr, "%s: no support for compressed BMP\n", path_in);
				return 4;
		}

		data_start = *((int32_t *) (buf + 0xa));
		image_size = *((uint32_t *) (buf + 0x22));

		output = malloc(0x6 + 0x28 + image_size);
		output[0] = 'B'; // Header
		output[1] = 'M';
		*((uint32_t *) (output + 0x2)) = 0x6 + 0x28 + image_size; // File size
		*((uint32_t *) (output + 0xa)) = 0x36; // Data start
		*((uint32_t *) (output + 0xe)) = 0x28; // Header size
		*((uint32_t *) (output + 0x12)) = *((uint32_t *) (buf + 0x12)); // Width
		*((uint32_t *) (output + 0x16)) = *((uint32_t *) (buf + 0x16)); // Height
		*((uint16_t *) (output + 0x1a)) = 1; // nb plan
		*((uint16_t *) (output + 0x1c)) = 24; // bpp
		*((uint32_t *) (output + 0x1e)) = 0; // compression
		*((uint32_t *) (output + 0x22)) = image_size; // Image size

		// Here's the vuln.
		memcpy(output + 0x36, buf + data_start, image_size);

		data_len = strlen(data);
		for (i = 0; i < image_size - 8 && i < data_len; i += 8) {
				output[0x36 + i + 0] = output[0x36 + i + 0] & 254 | ((data[i] & (1 << 7)) >> 7);
				output[0x36 + i + 1] = output[0x36 + i + 1] & 254 | ((data[i] & (1 << 6)) >> 6);
				output[0x36 + i + 2] = output[0x36 + i + 2] & 254 | ((data[i] & (1 << 5)) >> 5);
				output[0x36 + i + 3] = output[0x36 + i + 3] & 254 | ((data[i] & (1 << 4)) >> 4);
				output[0x36 + i + 4] = output[0x36 + i + 4] & 254 | ((data[i] & (1 << 3)) >> 3);
				output[0x36 + i + 5] = output[0x36 + i + 5] & 254 | ((data[i] & (1 << 2)) >> 2);
				output[0x36 + i + 6] = output[0x36 + i + 6] & 254 | ((data[i] & (1 << 1)) >> 1);
				output[0x36 + i + 7] = output[0x36 + i + 7] & 254 | ((data[i] & (1 << 0)) >> 0);
		}

		fp = fopen(path_out, "wb+");
		if (fp == NULL) {
				fprintf(stderr, "%s: could not write\n", path_out);
				return 5;
		}
		fwrite(output, 0x6 + 0x28 + image_size, 1, fp);
		fclose(fp);

		free(buf);
		return 0;
}