Blocktree/tools/rand_array.py

#!/usr/bin/env python3
'''
This script generates a rust array filled with random bytes.
'''

import random
import unittest
import sys

LINESEP = "\n"

def better_hex(byte: int) -> str:
    return "0x{:02X}".format(byte)

def make_array(bytes: bytes, col_limit: int = 100, indent: int = 4) -> str:
    indent_str = " " * (indent - 1)
    output = "[" + LINESEP
    line = indent_str
    for byte in bytes:
        next = " " + better_hex(byte) + ","
        if len(line) + len(next) >= col_limit:
            output += line + LINESEP
            line = indent_str
        line += next
    output += line + LINESEP + "];"
    return output

class TestMakeArray(unittest.TestCase):
    def test_make_array(self):
        bytes = b"\x8b!\xe9z\x14\xbc\xc6\xef\xf9(Y\xd5\xbd\xe8\xfb\x8a\x13I\xea,\xac\x9b\x7f\xf18\x8b\x94\x81\xe3\xf9!\xcc"
        expected = '''[
    0x8B, 0x21, 0xE9, 0x7A, 0x14, 0xBC, 0xC6, 0xEF, 0xF9, 0x28, 0x59, 0xD5, 0xBD, 0xE8, 0xFB, 0x8A,
    0x13, 0x49, 0xEA, 0x2C, 0xAC, 0x9B, 0x7F, 0xF1, 0x38, 0x8B, 0x94, 0x81, 0xE3, 0xF9, 0x21, 0xCC,
];'''
        actual = make_array(bytes)
        self.assertEqual(expected, actual)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        unittest.main()
    else:
        length = int(sys.argv[1])
        bytes = random.randbytes(length)
        array = make_array(bytes)
        print(array)