1649fd550341328077e403dd2b2024a9958ae2652Geremy Condra#! /usr/bin/env python
2649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
3649fd550341328077e403dd2b2024a9958ae2652Geremy Condraimport os
4649fd550341328077e403dd2b2024a9958ae2652Geremy Condraimport sys
5649fd550341328077e403dd2b2024a9958ae2652Geremy Condraimport struct
6649fd550341328077e403dd2b2024a9958ae2652Geremy Condraimport tempfile
7649fd550341328077e403dd2b2024a9958ae2652Geremy Condraimport commands
8649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
9649fd550341328077e403dd2b2024a9958ae2652Geremy CondraVERSION = 0
10649fd550341328077e403dd2b2024a9958ae2652Geremy CondraMAGIC_NUMBER = 0xb001b001
11649fd550341328077e403dd2b2024a9958ae2652Geremy CondraBLOCK_SIZE = 4096
12649fd550341328077e403dd2b2024a9958ae2652Geremy CondraMETADATA_SIZE = BLOCK_SIZE * 8
13649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
14649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef run(cmd):
15649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    status, output = commands.getstatusoutput(cmd)
16649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    print output
17649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    if status:
18649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        exit(-1)
19649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
20649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef get_verity_metadata_size(data_size):
21649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    return METADATA_SIZE
22649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
23649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef build_metadata_block(verity_table, signature):
24649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    table_len = len(verity_table)
25649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    block = struct.pack("II256sI", MAGIC_NUMBER, VERSION, signature, table_len)
26649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    block += verity_table
27649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    block = block.ljust(METADATA_SIZE, '\x00')
28649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    return block
29649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
30649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef sign_verity_table(table, signer_path, key_path):
31649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    with tempfile.NamedTemporaryFile(suffix='.table') as table_file:
32649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        with tempfile.NamedTemporaryFile(suffix='.sig') as signature_file:
33649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            table_file.write(table)
34649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            table_file.flush()
35649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            cmd = " ".join((signer_path, table_file.name, key_path, signature_file.name))
36649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            print cmd
37649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            run(cmd)
38649fd550341328077e403dd2b2024a9958ae2652Geremy Condra            return signature_file.read()
39649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
40649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef build_verity_table(block_device, data_blocks, root_hash, salt):
41649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    table = "1 %s %s %s %s %s %s sha256 %s %s"
42649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    table %= (  block_device,
43649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                block_device,
44649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                BLOCK_SIZE,
45649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                BLOCK_SIZE,
46649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                data_blocks,
4787757783008a561b9c54b274e03791e9dc76377fSami Tolvanen                data_blocks,
48649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                root_hash,
49649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                salt)
50649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    return table
51649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
52649fd550341328077e403dd2b2024a9958ae2652Geremy Condradef build_verity_metadata(data_blocks, metadata_image, root_hash,
53649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                            salt, block_device, signer_path, signing_key):
54649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    # build the verity table
55649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    verity_table = build_verity_table(block_device, data_blocks, root_hash, salt)
56649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    # build the verity table signature
57649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    signature = sign_verity_table(verity_table, signer_path, signing_key)
58649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    # build the metadata block
59649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    metadata_block = build_metadata_block(verity_table, signature)
60649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    # write it to the outfile
61649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    with open(metadata_image, "wb") as f:
62649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        f.write(metadata_block)
63649fd550341328077e403dd2b2024a9958ae2652Geremy Condra
64649fd550341328077e403dd2b2024a9958ae2652Geremy Condraif __name__ == "__main__":
65649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    if len(sys.argv) == 3 and sys.argv[1] == "-s":
66649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        print get_verity_metadata_size(int(sys.argv[2]))
67649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    elif len(sys.argv) == 8:
68649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        data_image_blocks = int(sys.argv[1]) / 4096
69649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        metadata_image = sys.argv[2]
70649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        root_hash = sys.argv[3]
71649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        salt = sys.argv[4]
72649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        block_device = sys.argv[5]
73649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        signer_path = sys.argv[6]
74649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        signing_key = sys.argv[7]
75649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        build_verity_metadata(data_image_blocks, metadata_image, root_hash,
76649fd550341328077e403dd2b2024a9958ae2652Geremy Condra                                salt, block_device, signer_path, signing_key)
77649fd550341328077e403dd2b2024a9958ae2652Geremy Condra    else:
78649fd550341328077e403dd2b2024a9958ae2652Geremy Condra        exit(-1)
79