1/*
2 *
3 *  Copyright (C) 2012 OpenWrt.org
4 *  Copyright (C) 2012 Mikko Hissa <mikko.hissa@uta.fi>
5 *
6 *  This program is free software; you can redistribute it and/or modify it
7 *  under the terms of the GNU General Public License version 2 as published
8 *  by the Free Software Foundation.
9 *
10 */
11
12#include <stdio.h>
13#include <stdlib.h>
14#include <string.h>
15#include <stdarg.h>
16#include <libgen.h>
17#include <errno.h>
18#include <arpa/inet.h>
19#include <unistd.h>
20#include "md5.h"
21
22#define HDR_LEN                 0x60
23#define BUF_SIZE                0x200
24#define VERSION_SIZE            0x10
25#define MD5_SIZE                0x10
26#define PAD_SIZE                0x20
27
28#define DEFAULT_BLOCK_SIZE      65535
29
30#define DEFAULT_HEAD_VALUE      0x0
31#define DEFAULT_VERSION         "123"
32#define DEFAULT_MAGIC           0x12345678
33
34typedef struct {
35	uint32_t head;
36	uint32_t vendor_id;
37	uint32_t product_id;
38	uint8_t  version[VERSION_SIZE];
39	uint32_t firmware_type;
40	uint32_t filesize;
41	uint32_t zero;
42	uint8_t  md5sum[MD5_SIZE];
43	uint8_t  pad[PAD_SIZE];
44	uint32_t chksum;
45	uint32_t magic;
46} img_header;
47
48typedef struct {
49	uint8_t id;
50	char * name;
51} firmware_type;
52
53typedef enum {
54	NONE, ENCODE, DECODE
55} op_mode;
56
57static firmware_type FIRMWARE_TYPES[] = {
58	{ 0x01, "bootloader" },
59	{ 0x02, "kernel" },
60	{ 0x03, "kernelapp" },
61	{ 0x04, "apps" },
62	/* The types below this line vary by manufacturer */
63	{ 0x05, "littleapps (D-Link)/factoryapps (EnGenius)" },
64	{ 0x06, "sounds (D-Link)/littleapps (EnGenius)" },
65	{ 0x07, "userconfig (D-Link)/appdata (EnGenius)" },
66	{ 0x08, "userconfig (EnGenius)"},
67	{ 0x09, "odmapps (EnGenius)"},
68	{ 0x0a, "factoryapps (D-Link)" },
69	{ 0x0b, "odmapps (D-Link)" },
70	{ 0x0c, "langpack (D-Link)" }
71};
72
73static long get_file_size(const char *filename)
74{
75	FILE *fp_file;
76	long result;
77
78	fp_file = fopen(filename, "r");
79	if (!fp_file)
80		return -1;
81	fseek(fp_file, 0, SEEK_END);
82	result = ftell(fp_file);
83	fclose(fp_file);
84	return result;
85}
86
87static int header_checksum(void *data, int len)
88{
89	int i;
90	int sum;
91
92	sum = 0;
93	if (data != NULL && len >= 0) {
94		for (i = 0; i < len; ++i)
95			sum += *(unsigned char *) (data + i);
96		return sum;
97	}
98
99	return -1;
100}
101
102static int md5_file(const char *filename, uint8_t *dst)
103{
104	FILE *fp_src;
105	MD5_CTX ctx;
106	char buf[BUF_SIZE];
107	size_t bytes_read;
108
109	MD5_Init(&ctx);
110
111	fp_src = fopen(filename, "r+b");
112	if (!fp_src) {
113		return -1;
114	}
115	while (!feof(fp_src)) {
116		bytes_read = fread(&buf, 1, BUF_SIZE, fp_src);
117		MD5_Update(&ctx, &buf, bytes_read);
118	}
119	fclose(fp_src);
120
121	MD5_Final(dst, &ctx);
122
123	return 0;
124}
125
126static int encode_image(const char *input_file_name,
127		const char *output_file_name, img_header *header, int block_size)
128{
129	char buf[BUF_SIZE];
130	size_t bytes_read;
131	size_t pad_len = 0;
132	size_t bytes_avail;
133
134	FILE *fp_input;
135	FILE *fp_output;
136
137	int i;
138	long magic;
139
140	fp_input = fopen(input_file_name, "r+b");
141	if (!fp_input) {
142		fprintf(stderr, "Cannot open %s !!\n", input_file_name);
143		return -1;
144	}
145
146	fp_output = fopen(output_file_name, "w+b");
147	if (!fp_output) {
148		fprintf(stderr, "Cannot open %s !!\n", output_file_name);
149		fclose(fp_input);
150		return -1;
151	}
152
153	header->filesize = get_file_size(input_file_name);
154	if (!header->filesize) {
155		fprintf(stderr, "File %s open/size error!\n", input_file_name);
156		fclose(fp_input);
157		fclose(fp_output);
158		return -1;
159	}
160	/*
161	 * Zero padding
162	 */
163	if (block_size > 0) {
164		pad_len = block_size - (header->filesize % block_size);
165	}
166
167	if (md5_file(input_file_name, (uint8_t *) &header->md5sum) < 0) {
168		fprintf(stderr, "MD5 failed on file %s\n", input_file_name);
169		fclose(fp_input);
170		fclose(fp_output);
171		return -1;
172	}
173	header->zero = 0;
174	header->chksum = header_checksum(header, HDR_LEN);
175	header->head = htonl(header->head);
176	header->vendor_id = htonl(header->vendor_id);
177	header->product_id = htonl(header->product_id);
178	header->firmware_type = htonl(header->firmware_type);
179	header->filesize = htonl(header->filesize);
180	header->chksum = htonl(header->chksum);
181	magic = header->magic;
182	header->magic = htonl(header->magic);
183
184	fwrite(header, HDR_LEN, 1, fp_output);
185
186	while (!feof(fp_input) || pad_len > 0) {
187
188		if (!feof(fp_input))
189			bytes_read = fread(&buf, 1, BUF_SIZE, fp_input);
190		else
191			bytes_read = 0;
192
193		/*
194		 * No more bytes read, start padding
195		 */
196		if (bytes_read < BUF_SIZE && pad_len > 0) {
197			bytes_avail = BUF_SIZE - bytes_read;
198			memset( &buf[bytes_read], 0, bytes_avail);
199			bytes_read += bytes_avail < pad_len ? bytes_avail : pad_len;
200			pad_len -= bytes_avail < pad_len ? bytes_avail : pad_len;
201		}
202
203		for (i = 0; i < bytes_read; i++)
204			buf[i] ^= magic >> (i % 8) & 0xff;
205		fwrite(&buf, bytes_read, 1, fp_output);
206	}
207
208	fclose(fp_input);
209	fclose(fp_output);
210	return 1;
211}
212
213int decode_image(const char *input_file_name, const char *output_file_name)
214{
215	img_header header;
216	char buf[BUF_SIZE];
217
218	FILE *fp_input;
219	FILE *fp_output;
220	unsigned int i;
221
222	size_t bytes_read;
223	size_t bytes_written;
224
225	fp_input = fopen(input_file_name, "r+b");
226	if (!fp_input) {
227		fprintf(stderr, "Cannot open %s !!\n", input_file_name);
228		fclose(fp_input);
229		return -1;
230	}
231
232	fp_output = fopen(output_file_name, "w+b");
233	if (!fp_output) {
234		fprintf(stderr, "Cannot open %s !!\n", output_file_name);
235		fclose(fp_output);
236		return -1;
237	}
238
239	if (fread(&header, 1, HDR_LEN, fp_input) != HDR_LEN) {
240		fprintf(stderr, "Incorrect header size!!");
241		fclose(fp_input);
242		fclose(fp_output);
243		return -1;
244	}
245
246	header.head = ntohl(header.head);
247	header.vendor_id = ntohl(header.vendor_id);
248	header.product_id = ntohl(header.product_id);
249	header.firmware_type = ntohl(header.firmware_type);
250	header.filesize = ntohl(header.filesize);
251	header.chksum = ntohl(header.chksum);
252	header.magic = ntohl(header.magic);
253
254	bytes_written = 0;
255	while (!feof(fp_input)) {
256
257		bytes_read = fread(&buf, 1, BUF_SIZE, fp_input);
258		for (i = 0; i < bytes_read; i++)
259			buf[i] ^= header.magic >> (i % 8) & 0xff;
260
261		/*
262		 * Handle padded source file
263		 */
264		if (bytes_written + bytes_read > header.filesize) {
265			bytes_read = header.filesize - bytes_written;
266			if (bytes_read > 0)
267				fwrite(&buf, bytes_read, 1, fp_output);
268			break;
269		}
270
271		fwrite(&buf, bytes_read, 1, fp_output);
272		bytes_written += bytes_read;
273	}
274
275	fclose(fp_input);
276	fclose(fp_output);
277
278	return 1;
279}
280
281static void usage(const char *progname, int status)
282{
283	FILE *stream = (status != EXIT_SUCCESS) ? stderr : stdout;
284	int i;
285
286	fprintf(stream, "Usage: %s [OPTIONS...]\n", progname);
287	fprintf(stream, "\n"
288			"Options:\n"
289			"  -e <file>		encode image file <file>\n"
290			"  -d <file>		decode image file <file>\n"
291			"  -o <file>		write output to the file <file>\n"
292			"  -t <type>		set image type to <type>\n"
293			"			valid image <type> values:\n");
294	for (i = 0; i < sizeof(FIRMWARE_TYPES) / sizeof(firmware_type); i++) {
295		fprintf(stream, "			%-5i= %s\n", FIRMWARE_TYPES[i].id,
296				FIRMWARE_TYPES[i].name);
297	}
298	fprintf(stream, "  -v <version>		set image version to <version>\n"
299			"  -r <vendor>		set image vendor id to <vendor>\n"
300			"  -p <product>		set image product id to <product>\n"
301			"  -m <magic>		set encoding magic <magic>\n"
302			"  -z			enable image padding to <blocksize>\n"
303			"  -b <blocksize>	set image <blocksize>, defaults to %u\n"
304			"  -h			show this screen\n", DEFAULT_BLOCK_SIZE);
305	exit(status);
306}
307
308int main(int argc, char *argv[])
309{
310	int opt;
311	char *input_file, *output_file, *progname = NULL;
312	op_mode mode = NONE;
313	int tmp, i, pad = 0;
314	int block_size;
315	img_header header;
316
317	block_size = DEFAULT_BLOCK_SIZE;
318	progname = basename(argv[0]);
319
320	memset(&header, 0, sizeof( img_header ));
321	header.magic = DEFAULT_MAGIC;
322	header.head = DEFAULT_HEAD_VALUE;
323	strncpy( (char*)&header.version, DEFAULT_VERSION, VERSION_SIZE - 1);
324
325	while ((opt = getopt(argc, argv, ":o:e:d:t:v:r:p:m:b:h?z")) != -1) {
326		switch (opt) {
327		case 'e':
328			input_file = optarg;
329			mode = ENCODE;
330			break;
331		case 'd':
332			input_file = optarg;
333			mode = DECODE;
334			break;
335		case 'o':
336			output_file = optarg;
337			break;
338		case 't':
339			tmp = strtol(optarg, 0, 10);
340			for (i = 0; i < sizeof(FIRMWARE_TYPES) / sizeof(firmware_type);
341					i++) {
342				if (FIRMWARE_TYPES[i].id == tmp) {
343					header.firmware_type = FIRMWARE_TYPES[i].id;
344					break;
345				}
346			}
347			if (header.firmware_type == 0) {
348				fprintf(stderr, "Invalid firmware type \"0\"!\n");
349				usage(progname, EXIT_FAILURE);
350			}
351			break;
352		case 'v':
353			strncpy( (char*)&header.version, optarg,
354					VERSION_SIZE - 1);
355			break;
356		case 'r':
357			header.vendor_id = strtol(optarg, 0, 0);
358			break;
359		case 'p':
360			header.product_id = strtol(optarg, 0, 0);
361			break;
362		case 'm':
363			header.magic = strtoul(optarg, 0, 16);
364			break;
365		case 'z':
366			pad = 1;
367			break;
368		case 'b':
369			block_size = strtol(optarg, 0, 10);
370			break;
371		case 'h':
372			usage(progname, EXIT_SUCCESS);
373			break;
374		case ':':
375			fprintf(stderr, "Option -%c requires an operand\n", optopt);
376			usage(progname, EXIT_FAILURE);
377			break;
378		case '?':
379			fprintf(stderr, "Unrecognized option: -%c\n", optopt);
380			usage(progname, EXIT_FAILURE);
381			break;
382		default:
383			usage(progname, EXIT_FAILURE);
384			break;
385		}
386	}
387
388	/* Check required arguments*/
389	if (header.firmware_type == 0) {
390		fprintf(stderr, "Firmware type must be defined\n");
391		usage(progname, EXIT_FAILURE);
392	} else if (input_file == 0 || output_file == 0) {
393		fprintf(stderr, "Input and output files must be defined\n");
394		usage(progname, EXIT_FAILURE);
395	} else if (header.vendor_id == 0 || header.product_id == 0) {
396		fprintf(stderr, "Vendor ID and Product ID must be defined and non-zero\n");
397		usage(progname, EXIT_FAILURE);
398	}
399
400	switch (mode) {
401	case NONE:
402		fprintf(stderr, "A mode must be defined\n");
403		usage(progname, EXIT_FAILURE);
404		break;
405	case ENCODE:
406		if (encode_image(input_file, output_file, &header, pad ? block_size : 0)
407				< 0)
408			return EXIT_FAILURE;
409		break;
410	case DECODE:
411		if (decode_image(input_file, output_file) < 0)
412			return EXIT_FAILURE;
413		break;
414	}
415
416	return EXIT_SUCCESS;
417}
418