message-capture-parser.py
1 #!/usr/bin/env python3 2 # Copyright (c) 2020-present The Bitcoin Core developers 3 # Distributed under the MIT software license, see the accompanying 4 # file COPYING or http://www.opensource.org/licenses/mit-license.php. 5 """Parse message capture binary files. To be used in conjunction with -capturemessages.""" 6 7 import argparse 8 import os 9 import shutil 10 import sys 11 from io import BytesIO 12 import json 13 from pathlib import Path 14 from typing import Any, Optional 15 16 sys.path.append(os.path.join(os.path.dirname(__file__), '../../test/functional')) 17 18 from test_framework.messages import ser_uint256 # noqa: E402 19 from test_framework.p2p import MESSAGEMAP # noqa: E402 20 21 TIME_SIZE = 8 22 LENGTH_SIZE = 4 23 MSGTYPE_SIZE = 12 24 25 # The test framework classes stores hashes as large ints in many cases. 26 # These are variables of type uint256 in core. 27 # There isn't a way to distinguish between a large int and a large int that is actually a blob of bytes. 28 # As such, they are itemized here. 29 # Any variables with these names that are of type int are actually uint256 variables. 30 # (These can be easily found by looking for calls to deser_uint256, deser_uint256_vector, and uint256_from_str in messages.py) 31 HASH_INTS = [ 32 "blockhash", 33 "block_hash", 34 "hash", 35 "hashMerkleRoot", 36 "hashPrevBlock", 37 "hashstop", 38 "prev_header", 39 "sha256", 40 "stop_hash", 41 ] 42 43 HASH_INT_VECTORS = [ 44 "hashes", 45 "headers", 46 "vHave", 47 "vHash", 48 ] 49 50 51 class ProgressBar: 52 def __init__(self, total: float): 53 self.total = total 54 self.running = 0 55 56 def set_progress(self, progress: float): 57 cols = shutil.get_terminal_size()[0] 58 if cols <= 12: 59 return 60 max_blocks = cols - 9 61 num_blocks = int(max_blocks * progress) 62 print('\r[ {}{} ] {:3.0f}%' 63 .format('#' * num_blocks, 64 ' ' * (max_blocks - num_blocks), 65 progress * 100), 66 end ='') 67 68 def update(self, more: float): 69 self.running += more 70 self.set_progress(self.running / self.total) 71 72 73 def to_jsonable(obj: Any) -> Any: 74 if hasattr(obj, "__dict__"): 75 return obj.__dict__ 76 elif hasattr(obj, "__slots__"): 77 ret = {} # type: Any 78 for slot in obj.__slots__: 79 val = getattr(obj, slot, None) 80 if slot in HASH_INTS and isinstance(val, int): 81 ret[slot] = ser_uint256(val).hex() 82 elif slot in HASH_INT_VECTORS and all(isinstance(a, int) for a in val): 83 ret[slot] = [ser_uint256(a).hex() for a in val] 84 else: 85 ret[slot] = to_jsonable(val) 86 return ret 87 elif isinstance(obj, list): 88 return [to_jsonable(a) for a in obj] 89 elif isinstance(obj, bytes): 90 return obj.hex() 91 else: 92 return obj 93 94 95 def process_file(path: str, messages: list[Any], recv: bool, progress_bar: Optional[ProgressBar]) -> None: 96 with open(path, 'rb') as f_in: 97 if progress_bar: 98 bytes_read = 0 99 100 while True: 101 if progress_bar: 102 # Update progress bar 103 diff = f_in.tell() - bytes_read - 1 104 progress_bar.update(diff) 105 bytes_read = f_in.tell() - 1 106 107 # Read the Header 108 tmp_header_raw = f_in.read(TIME_SIZE + LENGTH_SIZE + MSGTYPE_SIZE) 109 if not tmp_header_raw: 110 break 111 tmp_header = BytesIO(tmp_header_raw) 112 time = int.from_bytes(tmp_header.read(TIME_SIZE), "little") # type: int 113 msgtype = tmp_header.read(MSGTYPE_SIZE).split(b'\x00', 1)[0] # type: bytes 114 length = int.from_bytes(tmp_header.read(LENGTH_SIZE), "little") # type: int 115 116 # Start converting the message to a dictionary 117 msg_dict = {} 118 msg_dict["direction"] = "recv" if recv else "sent" 119 msg_dict["time"] = time 120 msg_dict["size"] = length # "size" is less readable here, but more readable in the output 121 122 msg_ser = BytesIO(f_in.read(length)) 123 124 # Determine message type 125 if msgtype not in MESSAGEMAP: 126 # Unrecognized message type 127 try: 128 msgtype_tmp = msgtype.decode() 129 if not msgtype_tmp.isprintable(): 130 raise UnicodeDecodeError 131 msg_dict["msgtype"] = msgtype_tmp 132 except UnicodeDecodeError: 133 msg_dict["msgtype"] = "UNREADABLE" 134 msg_dict["body"] = msg_ser.read().hex() 135 msg_dict["error"] = "Unrecognized message type." 136 messages.append(msg_dict) 137 print(f"WARNING - Unrecognized message type {msgtype} in {path}", file=sys.stderr) 138 continue 139 140 # Deserialize the message 141 msg = MESSAGEMAP[msgtype]() 142 msg_dict["msgtype"] = msgtype.decode() 143 144 try: 145 msg.deserialize(msg_ser) 146 except KeyboardInterrupt: 147 raise 148 except Exception: 149 # Unable to deserialize message body 150 msg_ser.seek(0, os.SEEK_SET) 151 msg_dict["body"] = msg_ser.read().hex() 152 msg_dict["error"] = "Unable to deserialize message." 153 messages.append(msg_dict) 154 print(f"WARNING - Unable to deserialize message in {path}", file=sys.stderr) 155 continue 156 157 # Convert body of message into a jsonable object 158 if length: 159 msg_dict["body"] = to_jsonable(msg) 160 messages.append(msg_dict) 161 162 if progress_bar: 163 # Update the progress bar to the end of the current file 164 # in case we exited the loop early 165 f_in.seek(0, os.SEEK_END) # Go to end of file 166 diff = f_in.tell() - bytes_read - 1 167 progress_bar.update(diff) 168 169 170 def main(): 171 parser = argparse.ArgumentParser( 172 description=__doc__, 173 epilog="EXAMPLE \n\t{0} -o out.json <data-dir>/message_capture/**/*.dat".format(sys.argv[0]), 174 formatter_class=argparse.RawTextHelpFormatter) 175 parser.add_argument( 176 "capturepaths", 177 nargs='+', 178 help="binary message capture files to parse.") 179 parser.add_argument( 180 "-o", "--output", 181 help="output file. If unset print to stdout") 182 parser.add_argument( 183 "-n", "--no-progress-bar", 184 action='store_true', 185 help="disable the progress bar. Automatically set if the output is not a terminal") 186 args = parser.parse_args() 187 capturepaths = [Path.cwd() / Path(capturepath) for capturepath in args.capturepaths] 188 output = Path.cwd() / Path(args.output) if args.output else False 189 use_progress_bar = (not args.no_progress_bar) and sys.stdout.isatty() 190 191 messages = [] # type: list[Any] 192 if use_progress_bar: 193 total_size = sum(capture.stat().st_size for capture in capturepaths) 194 progress_bar = ProgressBar(total_size) 195 else: 196 progress_bar = None 197 198 for capture in capturepaths: 199 process_file(str(capture), messages, "recv" in capture.stem, progress_bar) 200 201 messages.sort(key=lambda msg: msg['time']) 202 203 if use_progress_bar: 204 progress_bar.set_progress(1) 205 206 jsonrep = json.dumps(messages) 207 if output: 208 with open(str(output), 'w+') as f_out: 209 f_out.write(jsonrep) 210 else: 211 print(jsonrep) 212 213 if __name__ == "__main__": 214 main()