1/*
2 *  Copyright (C) 2011 Gabor Juhos <juhosg@openwrt.org>
3 *
4 *  This program is free software; you can redistribute it and/or modify it
5 *  under the terms of the GNU General Public License version 2 as published
6 *  by the Free Software Foundation.
7 *
8 */
9
10#include <stdio.h>
11#include <stdlib.h>
12#include <stdint.h>
13#include <string.h>
14#include <unistd.h>
15#include <libgen.h>
16#include <getopt.h>
17#include <stdarg.h>
18#include <errno.h>
19#include <sys/stat.h>
20
21#include "md5.h"
22
23#define ERR(fmt, ...) do { \
24	fflush(0); \
25	fprintf(stderr, "[%s] *** error: " fmt "\n", \
26			progname, ## __VA_ARGS__ ); \
27} while (0)
28
29#define ERRS(fmt, ...) do { \
30	int save = errno; \
31	fflush(0); \
32	fprintf(stderr, "[%s] *** error: " fmt ", %s\n", \
33			progname, ## __VA_ARGS__, strerror(save)); \
34} while (0)
35
36#define WRG_MAGIC	0x20040220
37
38struct wrg_header {
39	char		signature[32];
40	uint32_t	magic1;
41	uint32_t	magic2;
42	uint32_t	size;
43	uint32_t	offset;
44	char		devname[32];
45	char		digest[16];
46} __attribute__ ((packed));
47
48static char *progname;
49static char *ifname;
50static char *ofname;
51static char *signature;
52static char *dev_name;
53static uint32_t offset;
54static int big_endian;
55
56void usage(int status)
57{
58	FILE *stream = (status != EXIT_SUCCESS) ? stderr : stdout;
59
60	fprintf(stream, "Usage: %s [OPTIONS...]\n", progname);
61	fprintf(stream,
62"\n"
63"Options:\n"
64"  -b              create image in big endian format\n"
65"  -i <file>       read input from the file <file>\n"
66"  -d <name>       set device name to <name>\n"
67"  -o <file>       write output to the file <file>\n"
68"  -O <offset>     set offset to <offset>\n"
69"  -s <sig>        set image signature to <sig>\n"
70"  -h              show this screen\n"
71	);
72
73	exit(status);
74}
75
76static void put_u32(void *data, uint32_t val)
77{
78	unsigned char *p = data;
79
80	if (big_endian) {
81		p[0] = (val >> 24) & 0xff;
82		p[1] = (val >> 16) & 0xff;
83		p[2] = (val >> 8) & 0xff;
84		p[3] = val & 0xff;
85	} else {
86		p[3] = (val >> 24) & 0xff;
87		p[2] = (val >> 16) & 0xff;
88		p[1] = (val >> 8) & 0xff;
89		p[0] = val & 0xff;
90	}
91}
92
93static void get_digest(struct wrg_header *header, char *data, int size)
94{
95	MD5_CTX ctx;
96
97	MD5_Init(&ctx);
98
99	MD5_Update(&ctx, (char *)&header->offset, sizeof(header->offset));
100	MD5_Update(&ctx, (char *)&header->devname, sizeof(header->devname));
101	MD5_Update(&ctx, data, size);
102
103	MD5_Final(header->digest, &ctx);
104}
105
106int main(int argc, char *argv[])
107{
108	struct wrg_header *header;
109	char *buf;
110	struct stat st;
111	int buflen;
112	int err;
113	int res = EXIT_FAILURE;
114
115	FILE *outfile, *infile;
116
117	progname = basename(argv[0]);
118
119	while ( 1 ) {
120		int c;
121
122		c = getopt(argc, argv, "bd:i:o:s:O:h");
123		if (c == -1)
124			break;
125
126		switch (c) {
127		case 'b':
128			big_endian = 1;
129			break;
130		case 'd':
131			dev_name = optarg;
132			break;
133		case 'i':
134			ifname = optarg;
135			break;
136		case 'o':
137			ofname = optarg;
138			break;
139		case 's':
140			signature = optarg;
141			break;
142		case 'O':
143			offset = strtoul(optarg, NULL, 0);
144			break;
145		case 'h':
146			usage(EXIT_SUCCESS);
147			break;
148
149		default:
150			usage(EXIT_FAILURE);
151			break;
152		}
153	}
154
155	if (signature == NULL) {
156		ERR("no signature specified");
157		goto err;
158	}
159
160	if (ifname == NULL) {
161		ERR("no input file specified");
162		goto err;
163	}
164
165	if (ofname == NULL) {
166		ERR("no output file specified");
167		goto err;
168	}
169
170	if (dev_name == NULL) {
171		ERR("no device name specified");
172		goto err;
173	}
174
175	err = stat(ifname, &st);
176	if (err){
177		ERRS("stat failed on %s", ifname);
178		goto err;
179	}
180
181	buflen = st.st_size + sizeof(struct wrg_header);
182	buf = malloc(buflen);
183	if (!buf) {
184		ERR("no memory for buffer\n");
185		goto err;
186	}
187
188	infile = fopen(ifname, "r");
189	if (infile == NULL) {
190		ERRS("could not open \"%s\" for reading", ifname);
191		goto err_free;
192	}
193
194	errno = 0;
195	fread(buf + sizeof(struct wrg_header), st.st_size, 1, infile);
196	if (errno != 0) {
197		ERRS("unable to read from file %s", ifname);
198		goto close_in;
199	}
200
201	header = (struct wrg_header *) buf;
202	memset(header, '\0', sizeof(struct wrg_header));
203
204	strncpy(header->signature, signature, sizeof(header->signature));
205	strncpy(header->devname, dev_name, sizeof(header->signature));
206	put_u32(&header->magic1, WRG_MAGIC);
207	put_u32(&header->magic2, WRG_MAGIC);
208	put_u32(&header->size, st.st_size);
209	put_u32(&header->offset, offset);
210
211	get_digest(header, buf + sizeof(struct wrg_header), st.st_size);
212
213	outfile = fopen(ofname, "w");
214	if (outfile == NULL) {
215		ERRS("could not open \"%s\" for writing", ofname);
216		goto close_in;
217	}
218
219	errno = 0;
220	fwrite(buf, buflen, 1, outfile);
221	if (errno) {
222		ERRS("unable to write to file %s", ofname);
223		goto close_out;
224	}
225
226	fflush(outfile);
227
228	res = EXIT_SUCCESS;
229
230close_out:
231	fclose(outfile);
232	if (res != EXIT_SUCCESS)
233		unlink(ofname);
234close_in:
235	fclose(infile);
236err_free:
237	free(buf);
238err:
239	return res;
240}
241