1#!/usr/bin/env python3
2# Copyright (c) 2020 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
7import argparse
8import os
9import shutil
10import sys
11from io import BytesIO
12import json
13from pathlib import Path
14from typing import Any, List, Optional
15
16sys.path.append(os.path.join(os.path.dirname(__file__), '../../test/functional'))
17
18from test_framework.messages import ser_uint256     # noqa: E402
19from test_framework.p2p import MESSAGEMAP           # noqa: E402
20
21TIME_SIZE = 8
22LENGTH_SIZE = 4
23MSGTYPE_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)
31HASH_INTS = [
32    "blockhash",
33    "block_hash",
34    "hash",
35    "hashMerkleRoot",
36    "hashPrevBlock",
37    "hashstop",
38    "prev_header",
39    "sha256",
40    "stop_hash",
41]
42
43HASH_INT_VECTORS = [
44    "hashes",
45    "headers",
46    "vHave",
47    "vHash",
48]
49
50
51class 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
73def 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 isinstance(val[0], int):
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
95def 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
170def 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+', encoding="utf8") as f_out:
209            f_out.write(jsonrep)
210    else:
211        print(jsonrep)
212
213if __name__ == "__main__":
214    main()
215