1#! /usr/bin/env python
2
3import os
4import sys
5import struct
6import tempfile
7import commands
8
9VERSION = 0
10MAGIC_NUMBER = 0xb001b001
11BLOCK_SIZE = 4096
12METADATA_SIZE = BLOCK_SIZE * 8
13
14def run(cmd):
15    status, output = commands.getstatusoutput(cmd)
16    print output
17    if status:
18        exit(-1)
19
20def get_verity_metadata_size(data_size):
21    return METADATA_SIZE
22
23def build_metadata_block(verity_table, signature):
24    table_len = len(verity_table)
25    block = struct.pack("II256sI", MAGIC_NUMBER, VERSION, signature, table_len)
26    block += verity_table
27    block = block.ljust(METADATA_SIZE, '\x00')
28    return block
29
30def sign_verity_table(table, signer_path, key_path):
31    with tempfile.NamedTemporaryFile(suffix='.table') as table_file:
32        with tempfile.NamedTemporaryFile(suffix='.sig') as signature_file:
33            table_file.write(table)
34            table_file.flush()
35            cmd = " ".join((signer_path, table_file.name, key_path, signature_file.name))
36            print cmd
37            run(cmd)
38            return signature_file.read()
39
40def build_verity_table(block_device, data_blocks, root_hash, salt):
41    table = "1 %s %s %s %s %s %s sha256 %s %s"
42    table %= (  block_device,
43                block_device,
44                BLOCK_SIZE,
45                BLOCK_SIZE,
46                data_blocks,
47                data_blocks + (METADATA_SIZE / BLOCK_SIZE),
48                root_hash,
49                salt)
50    return table
51
52def build_verity_metadata(data_blocks, metadata_image, root_hash,
53                            salt, block_device, signer_path, signing_key):
54    # build the verity table
55    verity_table = build_verity_table(block_device, data_blocks, root_hash, salt)
56    # build the verity table signature
57    signature = sign_verity_table(verity_table, signer_path, signing_key)
58    # build the metadata block
59    metadata_block = build_metadata_block(verity_table, signature)
60    # write it to the outfile
61    with open(metadata_image, "wb") as f:
62        f.write(metadata_block)
63
64if __name__ == "__main__":
65    if len(sys.argv) == 3 and sys.argv[1] == "-s":
66        print get_verity_metadata_size(int(sys.argv[2]))
67    elif len(sys.argv) == 8:
68        data_image_blocks = int(sys.argv[1]) / 4096
69        metadata_image = sys.argv[2]
70        root_hash = sys.argv[3]
71        salt = sys.argv[4]
72        block_device = sys.argv[5]
73        signer_path = sys.argv[6]
74        signing_key = sys.argv[7]
75        build_verity_metadata(data_image_blocks, metadata_image, root_hash,
76                                salt, block_device, signer_path, signing_key)
77    else:
78        exit(-1)
79