1/*
2 * Copyright (C) 2016 The Android Open Source Project
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *      http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include <stdbool.h>
18#include <string.h>
19#include <stdint.h>
20#include <stdlib.h>
21#include <stdio.h>
22#include <inttypes.h>
23
24#include <nanohub/aes.h>
25#include <nanohub/sha2.h>
26#include <nanohub/nanohub.h>
27#include <nanohub/nanoapp.h>
28
29static FILE* urandom = NULL;
30
31static void cleanup(void)
32{
33    if (urandom)
34        fclose(urandom);
35}
36
37static void rand_bytes(void *dst, uint32_t len)
38{
39    if (!urandom) {
40        urandom = fopen("/dev/urandom", "rb");
41        if (!urandom) {
42            fprintf(stderr, "Failed to open /dev/urandom. Cannot procceed!\n");
43            exit(2);
44        }
45
46        //it might not matter, but we still like to try to cleanup after ourselves
47        (void)atexit(cleanup);
48    }
49
50    if (len != fread(dst, 1, len, urandom)) {
51        fprintf(stderr, "Failed to read /dev/urandom. Cannot procceed!\n");
52        exit(2);
53    }
54}
55
56static int handleEncrypt(uint8_t **pbuf, uint32_t bufUsed, FILE *out, uint64_t keyId, uint32_t *key)
57{
58    uint32_t i;
59    struct AesCbcContext ctx;
60    struct ImageHeader *image;
61    uint32_t *data;
62    struct Sha2state shaState;
63    bool err = false;
64    struct AppSecEncrHdr encr;
65    uint32_t padLen = 0;
66    uint8_t *buf = *pbuf;
67
68    encr.keyID = keyId;
69
70//FIXME: compatibility: all the devices has google secret key with id 1, so we
71//       can't simply change and enforce new key naming policy;
72//       first, key upload mechanism shall start working, and then we can have
73//       all the policies we want; for now, disable enforcement
74
75//        if (encr.keyID <= 0xFFFF)
76//            encr.keyID = AES_KEY_ID(encr.keyID);
77
78    fprintf(stderr, "Using Key ID: %016" PRIX64 "\n", encr.keyID);
79    rand_bytes(encr.IV, sizeof(encr.IV));
80    printHash(stderr, "Using IV", encr.IV, AES_BLOCK_WORDS);
81
82    if (bufUsed <= sizeof(*image)) {
83        fprintf(stderr, "Input file is too small\n");
84        return 2;
85    }
86
87    encr.dataLen = bufUsed;
88
89    if (((bufUsed - sizeof(*image)) % AES_BLOCK_SIZE) != 0)
90        padLen = AES_BLOCK_SIZE - ((bufUsed - sizeof(*image)) % AES_BLOCK_SIZE);
91
92    if (padLen) {
93        reallocOrDie(buf, bufUsed + padLen);
94        rand_bytes(buf + bufUsed, padLen);
95        bufUsed += padLen;
96        fprintf(stderr, "Padded to %" PRIu32 " bytes\n", bufUsed);
97        *pbuf = buf;
98    }
99
100    image = (struct ImageHeader *)buf;
101
102    if (bufUsed >= sizeof(*image) && image->aosp.magic == NANOAPP_AOSP_MAGIC &&
103        image->aosp.header_version == 1 && image->layout.magic == GOOGLE_LAYOUT_MAGIC) {
104        fprintf(stderr, "Found AOSP header\n");
105    } else {
106        fprintf(stderr, "Unknown binary format\n");
107        return 2;
108    }
109
110    if ((image->aosp.flags & NANOAPP_SIGNED_FLAG) != 0) {
111        fprintf(stderr, "data is marked as signed; encryption is not possible for signed data\n");
112        return 2;
113    }
114    if ((image->aosp.flags & NANOAPP_ENCRYPTED_FLAG) != 0) {
115        fprintf(stderr, "data is marked as encrypted; encryption is not possible for encrypted data\n");
116        return 2;
117    }
118
119    image->aosp.flags |= NANOAPP_ENCRYPTED_FLAG;
120    fwrite(image, sizeof(*image), 1, out);
121    data = (uint32_t *)(image + 1);
122    fprintf(stderr, "orig len: %" PRIu32 " bytes\n", encr.dataLen);
123    bufUsed -= sizeof(*image);
124    encr.dataLen -= sizeof(*image);
125    fwrite(&encr, sizeof(encr), 1, out);
126    sha2init(&shaState);
127
128    //encrypt and emit data
129    aesCbcInitForEncr(&ctx, key, encr.IV);
130    uint32_t outBuf[AES_BLOCK_WORDS];
131    for (i = 0; i < bufUsed/sizeof(uint32_t); i += AES_BLOCK_WORDS) {
132        aesCbcEncr(&ctx, data + i, outBuf);
133        int32_t sz = encr.dataLen - (i * sizeof(uint32_t));
134        sz = sz > AES_BLOCK_SIZE ? AES_BLOCK_SIZE : sz;
135        if (sz > 0) {
136            sha2processBytes(&shaState, data + i, sz);
137            fwrite(outBuf, AES_BLOCK_SIZE, 1, out);
138        }
139    }
140    const uint32_t *hash = sha2finish(&shaState);
141
142    printHash(stderr, "HASH", hash, SHA2_HASH_WORDS);
143
144    // finally, encrypt and output SHA2 hash
145    aesCbcEncr(&ctx, hash, outBuf);
146    fwrite(outBuf, AES_BLOCK_SIZE, 1, out);
147    aesCbcEncr(&ctx, hash + AES_BLOCK_WORDS, outBuf);
148    err = fwrite(outBuf, AES_BLOCK_SIZE, 1, out) != 1;
149
150    return err ? 2 : 0;
151}
152
153static int handleDecrypt(uint8_t **pbuf, uint32_t bufUsed, FILE *out, uint32_t *key)
154{
155    struct AesCbcContext ctx;
156    struct ImageHeader *image;
157    struct Sha2state shaState;
158    struct AppSecEncrHdr *encr;
159    uint32_t *data;
160    bool err = false;
161    uint32_t fileHash[((SHA2_HASH_WORDS + AES_BLOCK_WORDS - 1) / AES_BLOCK_WORDS) * AES_BLOCK_WORDS], fileHashSz;
162    uint32_t outBuf[AES_BLOCK_WORDS];
163    uint32_t i;
164    uint8_t *buf = *pbuf;
165
166    //parse header
167    image = (struct ImageHeader*)buf;
168    if (bufUsed >= (sizeof(*image) + sizeof(*encr)) &&
169        image->aosp.header_version == 1 && image->aosp.magic == NANOAPP_AOSP_MAGIC &&
170        image->layout.magic == GOOGLE_LAYOUT_MAGIC) {
171        fprintf(stderr, "Found AOSP header\n");
172        if (!(image->aosp.flags & NANOAPP_ENCRYPTED_FLAG)) {
173            fprintf(stderr, "data is not marked as encrypted; can't decrypt\n");
174            return 2;
175        }
176        image->aosp.flags &= ~NANOAPP_ENCRYPTED_FLAG;
177        data = (uint32_t *)(image + 1);
178        encr = (struct AppSecEncrHdr *)data;
179        data = (uint32_t *)(encr + 1);
180        bufUsed -= sizeof(*image) + sizeof(*encr);
181    } else {
182        fprintf(stderr, "Unknown binary format\n");
183        return 2;
184    }
185
186    if (encr->dataLen > bufUsed) {
187        fprintf(stderr, "Claimed output size of %" PRIu32 "b invalid\n", encr->dataLen);
188        return 2;
189    }
190    fprintf(stderr, "Original size %" PRIu32 "b (%" PRIu32 "b of padding present)\n",
191            encr->dataLen, bufUsed - encr->dataLen);
192    if (!encr->keyID)  {
193        fprintf(stderr, "Input data has invalid key ID\n");
194        return 2;
195    }
196    fprintf(stderr, "Using Key ID: %016" PRIX64 "\n", encr->keyID);
197    printHash(stderr, "Using IV", encr->IV, AES_BLOCK_WORDS);
198
199    fwrite(image, sizeof(*image), 1, out);
200        //decrypt and emit data
201    aesCbcInitForDecr(&ctx, key, encr->IV);
202    fileHashSz = 0;
203    sha2init(&shaState);
204    for (i = 0; i < bufUsed / sizeof(uint32_t); i += AES_BLOCK_WORDS) {
205        int32_t size = encr->dataLen - i * sizeof(uint32_t);
206        aesCbcDecr(&ctx, data + i, outBuf);
207        if (size > AES_BLOCK_SIZE)
208            size = AES_BLOCK_SIZE;
209        if (size > 0) {
210            sha2processBytes(&shaState, outBuf, size);
211            err = fwrite(outBuf, size, 1, out) != 1;
212        } else if (fileHashSz < sizeof(fileHash)) {
213            memcpy(((uint8_t*)fileHash) + fileHashSz, outBuf, AES_BLOCK_SIZE);
214            fileHashSz += AES_BLOCK_SIZE;
215        } else {
216            fprintf(stderr, "Too much input data\n");
217            return 2;
218        }
219    }
220    const uint32_t *calcHash = sha2finish(&shaState);
221    printHash(stderr, "HASH [calc]", calcHash, SHA2_HASH_WORDS);
222    printHash(stderr, "HASH [file]", fileHash, SHA2_HASH_WORDS);
223
224    bool verify = memcmp(fileHash, calcHash, SHA2_HASH_SIZE) == 0;
225    fprintf(stderr, "hash verification: %s\n", verify ? "passed" : "failed");
226    if (!verify)
227        return 2;
228
229    if (!err)
230        fprintf(stderr, "Done\n");
231
232    return err ? 2 : 0;
233}
234
235static void fatalUsage(const char *name, const char *msg, const char *arg)
236{
237    if (msg && arg)
238        fprintf(stderr, "Error: %s: %s\n\n", msg, arg);
239    else if (msg)
240        fprintf(stderr, "Error: %s\n\n", msg);
241
242    fprintf(stderr, "USAGE: %s [-e] [-d] [-i <key id>] [-k <key file>] <input file> [<output file>]\n"
243                    "       -i : 64-bit hex number != 0\n"
244                    "       -e : encrypt post-processed file\n"
245                    "       -d : decrypt encrypted post-processed file\n"
246                    "       -k : binary file (32 byte size) containing AES-256 secret key\n"
247                    , name);
248    exit(1);
249}
250
251int main(int argc, char **argv)
252{
253    uint32_t bufUsed = 0;
254    uint8_t *buf = NULL;
255    uint64_t keyId = 0;
256    int ret = -1;
257    uint32_t *u32Arg = NULL;
258    uint64_t *u64Arg = NULL;
259    const char **strArg = NULL;
260    const char *appName = argv[0];
261    const char *posArg[2] = { NULL };
262    uint32_t posArgCnt = 0;
263    FILE *out = NULL;
264    const char *prev = NULL;
265    bool decrypt = false;
266    bool encrypt = false;
267    const char *keyFile = NULL;
268    int multi = 0;
269    uint32_t key[AES_KEY_WORDS];
270
271    for (int i = 1; i < argc; i++) {
272        char *end = NULL;
273        if (argv[i][0] == '-') {
274            prev = argv[i];
275            if (!strcmp(argv[i], "-d"))
276                decrypt = true;
277            else if (!strcmp(argv[i], "-e"))
278                encrypt = true;
279            else if (!strcmp(argv[i], "-k"))
280                strArg = &keyFile;
281            else if (!strcmp(argv[i], "-i"))
282                u64Arg = &keyId;
283            else
284                fatalUsage(appName, "unknown argument", argv[i]);
285        } else {
286            if (u64Arg) {
287                uint64_t tmp = strtoull(argv[i], &end, 16);
288                if (*end == '\0')
289                    *u64Arg = tmp;
290                u64Arg = NULL;
291            } else if (u32Arg) {
292                uint32_t tmp = strtoul(argv[i], &end, 16);
293                if (*end == '\0')
294                    *u32Arg = tmp;
295                u32Arg = NULL;
296            } else if (strArg) {
297                    *strArg = argv[i];
298                strArg = NULL;
299            } else {
300                if (posArgCnt < 2)
301                    posArg[posArgCnt++] = argv[i];
302                else
303                    fatalUsage(appName, "too many positional arguments", argv[i]);
304            }
305            prev = 0;
306        }
307    }
308    if (prev)
309        fatalUsage(appName, "missing argument after", prev);
310
311    if (!posArgCnt)
312        fatalUsage(appName, "missing input file name", NULL);
313
314    if (encrypt)
315        multi++;
316    if (decrypt)
317        multi++;
318
319    if (multi != 1)
320        fatalUsage(appName, "select either -d or -e", NULL);
321
322    if (!keyFile)
323        fatalUsage(appName, "no key file given", NULL);
324
325    if (encrypt && !keyId)
326        fatalUsage(appName, "Non-zero Key ID must be given to encrypt data", NULL);
327
328    //read key
329    if (!readFile(key, sizeof(key), keyFile))
330        fatalUsage(appName, "Key file does not exist or has incorrect size", keyFile);
331
332    buf = loadFile(posArg[0], &bufUsed);
333    fprintf(stderr, "Read %" PRIu32 " bytes\n", bufUsed);
334
335    if (!posArg[1])
336        out = stdout;
337    else
338        out = fopen(posArg[1], "w");
339    if (!out)
340        fatalUsage(appName, "failed to create/open output file", posArg[1]);
341
342    if (encrypt)
343        ret = handleEncrypt(&buf, bufUsed, out, keyId, key);
344    else if (decrypt)
345        ret = handleDecrypt(&buf, bufUsed, out, key);
346
347    free(buf);
348    fclose(out);
349    return ret;
350}
351