71 lines
2.8 KiB
Python
71 lines
2.8 KiB
Python
"""Read a named embedded FileDescriptorProto without executing its carrier."""
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
from google.protobuf import descriptor_pb2, json_format
|
|
|
|
|
|
def varint(data, pos):
|
|
value = 0
|
|
for shift in range(0, 70, 7):
|
|
byte = data[pos]
|
|
pos += 1
|
|
value |= (byte & 127) << shift
|
|
if byte < 128:
|
|
return value, pos
|
|
raise ValueError('Oversized varint')
|
|
|
|
|
|
def extract(data, name):
|
|
encoded = name.encode()
|
|
occurrences = []
|
|
pos = 0
|
|
while (pos := data.find(encoded, pos)) >= 0:
|
|
if len(encoded) < 128 and data[pos-2:pos] == bytes([10, len(encoded)]):
|
|
occurrences.append(pos-2)
|
|
pos += len(encoded)
|
|
valid = []
|
|
for start in occurrences:
|
|
pos = start
|
|
while pos < min(len(data), start+16*1024*1024):
|
|
tag, pos = varint(data, pos)
|
|
if tag & 7 != 2 or tag >> 3 not in (1, 2, 3, 4, 5, 6, 7, 8, 9, 12, 14):
|
|
break
|
|
size, pos = varint(data, pos)
|
|
end = pos+size
|
|
if end > len(data):
|
|
break
|
|
if tag >> 3 == 12 and data[pos:end] in (b'proto2', b'proto3'):
|
|
desc = descriptor_pb2.FileDescriptorProto.FromString(data[start:end])
|
|
if desc.name == name:
|
|
valid.append((start, end, desc))
|
|
break
|
|
pos = end
|
|
if len(valid) != 1:
|
|
raise ValueError(f'Expected one complete descriptor {name}, found {len(valid)}')
|
|
return valid[0]
|
|
|
|
|
|
def main():
|
|
p = argparse.ArgumentParser(description=__doc__)
|
|
p.add_argument('--binary', type=Path, required=True)
|
|
p.add_argument('--descriptor', required=True)
|
|
p.add_argument('--output', type=Path, required=True)
|
|
args = p.parse_args()
|
|
data = args.binary.read_bytes()
|
|
start, end, desc = extract(data, args.descriptor)
|
|
args.output.mkdir(parents=True, exist_ok=True)
|
|
stem = Path(args.descriptor).stem
|
|
provenance = {'artifact_sha256':hashlib.sha256(data).hexdigest(), 'descriptor':desc.name, 'offset':start, 'length':end-start, 'descriptor_sha256':hashlib.sha256(data[start:end]).hexdigest()}
|
|
(args.output/(stem+'.descriptor.pb')).write_bytes(data[start:end])
|
|
(args.output/(stem+'.descriptor.json')).write_text(json.dumps({'provenance':provenance, 'descriptor':json_format.MessageToDict(desc, preserving_proto_field_name=True)}, indent=2), encoding='utf8')
|
|
cmds = {v.name:v.number for enum in desc.enum_type if enum.name.startswith('Cmd') for v in enum.value}
|
|
if cmds:
|
|
(args.output/'cmdids.json').write_text(json.dumps({'provenance':provenance, 'commands':cmds},indent=2), encoding='utf8')
|
|
print(json.dumps({**provenance, 'enums':len(desc.enum_type), 'messages':len(desc.message_type), 'commands':len(cmds)}))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|