commit 51382bb75c512fb0d48ced5e6f2fe4cdcc27fcb2
parent 443d11f56ef2f62783329ffef3e58be432372870
Author: nolash <dev@holbrook.no>
Date: Mon, 13 Sep 2021 22:26:23 +0200
Add stdin input to tool
Diffstat:
| M | test.py | | | 64 | +++++++++++++++++++++++++++++++++++++++++++++++++--------------- |
1 file changed, 49 insertions(+), 15 deletions(-)
diff --git a/test.py b/test.py
@@ -2,6 +2,7 @@
import argparse
import sys
import logging
+import select
logging.basicConfig(level=logging.WARNING)
logg = logging.getLogger()
@@ -11,29 +12,62 @@ argparser.add_argument('-n', action='store_true', help='skip newline at end of o
argparser.add_argument('-b', action='store_true', help='output raw bytes')
argparser.add_argument('-l', dest='data_length', type=int, help='length of data represented by preimage')
argparser.add_argument('-v', action='store_true', help='verbose output')
-argparser.add_argument('data', type=str, help='data input for BMT hasher')
+argparser.add_argument('-u', action='store_true', help='treat input as utf-8 string')
+argparser.add_argument('data', nargs='?', type=str, help='data input for BMT hasher')
largs = argparser.parse_args(sys.argv[1:])
+
+def stdin_arg():
+ """Retreive input arguments from stdin if they exist.
+
+ Method does not block, and expects arguments to be ready on stdin before being called.
+
+ :rtype: str
+ :returns: Input arguments string
+ """
+ h = select.select([sys.stdin.buffer], [], [])
+ if len(h[0]) > 0:
+ v = h[0][0].read()
+ #return v.rstrip()
+ return v
+ return None
+
+
if largs.v:
logg.setLevel(logging.DEBUG)
-input_data = largs.data.encode('utf-8')
+data = largs.data
+if data == None:
+ data = stdin_arg()
+
+input_data = data
+if not largs.u:
+ input_data.rstrip()
+ if isinstance(input_data, str):
+ input_data = data.encode('utf-8')
+
+input_data_length = 0
input_data_length = len(input_data)
data_length = input_data_length
if largs.data_length != None:
data_length = largs.data_length
logg.info('hashing {} bytes input with {} length prefix'.format(int(input_data_length), int(data_length)))
-logg.debug('data "{}"'.format(input_data.hex()))
-
-# TODO: why does this segfault just one line before?
-import swarm
-r = swarm.bmt(input_data, input_data_length, data_length)
-
-if largs.b:
- sys.stdout.buffer.write(r[:32])
-else:
- s = '{}'.format(r[:32].hex())
- if not largs.n:
- s += '\n'
- sys.stdout.write(s)
+
+
+def main():
+ # TODO: why does this segfault just one line before?
+ import swarm
+ r = swarm.bmt(input_data, input_data_length, data_length)
+
+ if largs.b:
+ sys.stdout.buffer.write(r[:32])
+ else:
+ s = '{}'.format(r[:32].hex())
+ if not largs.n:
+ s += '\n'
+ sys.stdout.write(s)
+
+
+if __name__ == '__main__':
+ main()
