1/*
2 *
3 *  Copyright (C) 2014 OpenWrt.org
4 *  Copyright (C) 2014 Mikko Hissa <mikko.hissa@werzek.com>
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 <errno.h>
13#include <fcntl.h>
14#include <stdio.h>
15#include <stdlib.h>
16#include <string.h>
17#include <netinet/in.h>
18#include <sys/mman.h>
19#include <sys/stat.h>
20#include <time.h>
21#include <unistd.h>
22#include <zlib.h>
23
24#define IH_MAGIC	0x27051956
25#define IH_NMLEN	32
26#define IH_PRODLEN	23
27
28#define IH_TYPE_INVALID		0
29#define IH_TYPE_STANDALONE	1
30#define IH_TYPE_KERNEL		2
31#define IH_TYPE_RAMDISK		3
32#define IH_TYPE_MULTI		4
33#define IH_TYPE_FIRMWARE	5
34#define IH_TYPE_SCRIPT		6
35#define IH_TYPE_FILESYSTEM	7
36
37/*
38 * Compression Types
39 */
40#define IH_COMP_NONE		0
41#define IH_COMP_GZIP		1
42#define IH_COMP_BZIP2		2
43#define IH_COMP_LZMA		3
44
45typedef struct {
46	uint8_t major;
47	uint8_t minor;
48} version_t;
49
50typedef struct {
51	version_t	kernel;
52	version_t	fs;
53	uint8_t		productid[IH_PRODLEN];
54	uint8_t  	sub_fs;
55	uint32_t	ih_ksz;
56} asus_t;
57
58typedef struct image_header {
59	uint32_t	ih_magic;
60	uint32_t	ih_hcrc;
61	uint32_t	ih_time;
62	uint32_t	ih_size;
63	uint32_t	ih_load;
64	uint32_t	ih_ep;
65	uint32_t	ih_dcrc;
66	uint8_t		ih_os;
67	uint8_t		ih_arch;
68	uint8_t		ih_type;
69	uint8_t		ih_comp;
70	union {
71		uint8_t	ih_name[IH_NMLEN];
72		asus_t	asus;
73	} tail;
74} image_header_t;
75
76typedef struct squashfs_sb {
77	uint32_t	s_magic;
78	uint32_t	pad0[9];
79	uint64_t	bytes_used;
80} squashfs_sb_t;
81
82typedef enum {
83	NONE, FACTORY, SYSUPGRADE,
84} op_mode_t;
85
86void
87calc_crc(image_header_t *hdr, void *data, uint32_t len)
88{
89	/*
90	 * Calculate payload checksum
91	 */
92	hdr->ih_dcrc = htonl(crc32(0, (Bytef *)data, len));
93	hdr->ih_size = htonl(len);
94	/*
95	 * Calculate header checksum
96	 */
97	hdr->ih_hcrc = 0;
98	hdr->ih_hcrc = htonl(crc32(0, (Bytef *)hdr, sizeof(image_header_t)));
99}
100
101
102static void
103usage(const char *progname, int status)
104{
105	FILE *stream = (status != EXIT_SUCCESS) ? stderr : stdout;
106	int i;
107
108	fprintf(stream, "Usage: %s [OPTIONS...]\n", progname);
109	fprintf(stream, "\n"
110			"Options:\n"
111			"  -f <file>		generate a factory flash image <file>\n"
112			"  -s <file>		generate a sysupgrade flash image <file>\n"
113			"  -h			show this screen\n");
114	exit(status);
115}
116
117int
118process_image(char *progname, char *filename, op_mode_t opmode)
119{
120	int 		fd, len;
121	void 		*data, *ptr;
122	char		namebuf[IH_NMLEN];
123	struct 		stat sbuf;
124	uint32_t	checksum, offset_kernel, offset_sqfs, offset_end,
125				offset_sec_header, offset_eb, offset_image_end;
126	squashfs_sb_t *sqs;
127	image_header_t *hdr;
128
129	if ((fd = open(filename, O_RDWR, 0666)) < 0) {
130		fprintf (stderr, "%s: Can't open %s: %s\n",
131			progname, filename, strerror(errno));
132		return (EXIT_FAILURE);
133	}
134
135	if (fstat(fd, &sbuf) < 0) {
136		fprintf (stderr, "%s: Can't stat %s: %s\n",
137			progname, filename, strerror(errno));
138		return (EXIT_FAILURE);
139	}
140
141	if ((unsigned)sbuf.st_size < sizeof(image_header_t)) {
142		fprintf (stderr,
143			"%s: Bad size: \"%s\" is no valid image\n",
144			progname, filename);
145		return (EXIT_FAILURE);
146	}
147
148	ptr = (void *)mmap(0, sbuf.st_size,
149				PROT_READ | PROT_WRITE,
150				MAP_SHARED,
151				fd, 0);
152
153	if ((caddr_t)ptr == (caddr_t)-1) {
154		fprintf (stderr, "%s: Can't read %s: %s\n",
155			progname, filename, strerror(errno));
156		return (EXIT_FAILURE);
157	}
158
159	hdr = ptr;
160
161	if (ntohl(hdr->ih_magic) != IH_MAGIC) {
162		fprintf (stderr,
163			"%s: Bad Magic Number: \"%s\" is no valid image\n",
164			progname, filename);
165		return (EXIT_FAILURE);
166	}
167
168	if (opmode == FACTORY) {
169		strncpy(namebuf, hdr->tail.ih_name, IH_NMLEN);
170		hdr->tail.asus.kernel.major = 0;
171		hdr->tail.asus.kernel.minor = 0;
172		hdr->tail.asus.fs.major = 0;
173		hdr->tail.asus.fs.minor = 0;
174		strncpy((char *)&hdr->tail.asus.productid, "RT-N56U", IH_PRODLEN);
175	}
176
177	if (hdr->tail.asus.ih_ksz == 0)
178		hdr->tail.asus.ih_ksz = htonl(ntohl(hdr->ih_size) + sizeof(image_header_t));
179
180	offset_kernel = sizeof(image_header_t);
181	offset_sqfs = ntohl(hdr->tail.asus.ih_ksz);
182	sqs = ptr + offset_sqfs;
183	offset_sec_header = offset_sqfs + sqs->bytes_used;
184
185	/*
186	 * Reserve space for the second header.
187	 */
188	offset_end = offset_sec_header + sizeof(image_header_t);
189	offset_eb = ((offset_end>>16)+1)<<16;
190
191	if (opmode == FACTORY)
192		offset_image_end = offset_eb + 4;
193	else
194		offset_image_end = sbuf.st_size;
195	/*
196	 * Move the second header at the end of the image.
197	 */
198	offset_end = offset_sec_header;
199	offset_sec_header = offset_eb - sizeof(image_header_t);
200
201	/*
202	 * Remove jffs2 markers between squashfs and eb boundary.
203	 */
204	if (opmode == FACTORY)
205		memset(ptr+offset_end, 0xff ,offset_eb - offset_end);
206
207	/*
208	 * Grow the image if needed.
209	 */
210	if (offset_image_end > sbuf.st_size) {
211		(void) munmap((void *)ptr, sbuf.st_size);
212		ftruncate(fd, offset_image_end);
213		ptr = (void *)mmap(0, offset_image_end,
214						PROT_READ | PROT_WRITE,
215						MAP_SHARED,
216						fd, 0);
217		/*
218		 * jffs2 marker
219		 */
220		if (opmode == FACTORY) {
221			*(uint8_t *)(ptr+offset_image_end-4) = 0xde;
222			*(uint8_t *)(ptr+offset_image_end-3) = 0xad;
223			*(uint8_t *)(ptr+offset_image_end-2) = 0xc0;
224			*(uint8_t *)(ptr+offset_image_end-1) = 0xde;
225		}
226	}
227
228	/*
229	 * Calculate checksums for the second header to be used after flashing.
230	 */
231	if (opmode == FACTORY) {
232		hdr = ptr+offset_sec_header;
233		memcpy(hdr, ptr, sizeof(image_header_t));
234		strncpy(hdr->tail.ih_name, namebuf, IH_NMLEN);
235		calc_crc(hdr, ptr+offset_kernel, offset_sqfs - offset_kernel);
236		calc_crc((image_header_t *)ptr, ptr+offset_kernel, offset_image_end - offset_kernel);
237	} else {
238		calc_crc((image_header_t *)ptr, ptr+offset_kernel, offset_sqfs - offset_kernel);
239	}
240
241	if (sbuf.st_size > offset_image_end)
242		(void) munmap((void *)ptr, sbuf.st_size);
243	else
244		(void) munmap((void *)ptr, offset_image_end);
245
246	ftruncate(fd, offset_image_end);
247	(void) close (fd);
248
249	return EXIT_SUCCESS;
250}
251
252int
253main(int argc, char **argv)
254{
255	int 		opt;
256	char 		*filename, *progname;
257	op_mode_t	opmode = NONE;
258
259	progname = argv[0];
260
261	while ((opt = getopt(argc, argv,":s:f:h?")) != -1) {
262		switch (opt) {
263		case 's':
264			opmode = SYSUPGRADE;
265			filename = optarg;
266			break;
267		case 'f':
268			opmode = FACTORY;
269			filename = optarg;
270			break;
271		case 'h':
272			opmode = NONE;
273		default:
274			usage(progname, EXIT_FAILURE);
275			opmode = NONE;
276		}
277	}
278
279	if(filename == NULL)
280		opmode = NONE;
281
282	switch (opmode) {
283	case NONE:
284		usage(progname, EXIT_FAILURE);
285		break;
286	case FACTORY:
287	case SYSUPGRADE:
288		return process_image(progname, filename, opmode);
289		break;
290	}
291
292	return EXIT_SUCCESS;
293}
294
295