Browse Source

util.py: improve bytespec-parsing algorithm

MMGen 6 years ago
parent
commit
a2386744e6
3 changed files with 20 additions and 8 deletions
  1. 1 1
      mmgen/seed.py
  2. 2 2
      mmgen/tool.py
  3. 17 5
      mmgen/util.py

+ 1 - 1
mmgen/seed.py

@@ -1022,7 +1022,7 @@ harder to find, you're advised to choose a much larger file size than this.
 				min_fsize = d.target_data_len + d.hincog_offset
 				msg(self.msg['choose_file_size'].format(min_fsize))
 				while True:
-					fsize = parse_nbytes(my_raw_input('Enter file size: '))
+					fsize = parse_bytespec(my_raw_input('Enter file size: '))
 					if fsize >= min_fsize: break
 					msg('File size must be an integer no less than {}'.format(min_fsize))
 

+ 2 - 2
mmgen/tool.py

@@ -476,7 +476,7 @@ class MMGenToolCmd(object):
 		return True
 
 	def rand2file(self,outfile,nbytes,threads=4,silent=False):
-		nbytes = parse_nbytes(nbytes)
+		nbytes = parse_bytespec(nbytes)
 		from Crypto import Random
 		rh = Random.new()
 		from queue import Queue
@@ -534,7 +534,7 @@ class MMGenToolCmd(object):
 		return True
 
 	def bytespec(self,s):
-		return str(parse_nbytes(s))
+		return str(parse_bytespec(s))
 
 	def keyaddrlist2monerowallets(self,infile,blockheight=None,addrs=None):
 		return self.monero_wallet_ops(infile=infile,op='create',blockheight=blockheight,addrs=addrs)

+ 17 - 5
mmgen/util.py

@@ -101,18 +101,30 @@ def set_for_type(val,refval,desc,invert_bool=False,src=None):
 # c=1, w=2, b=512, kB=1000, K=1024, MB=1000*1000, M=1024*1024,
 # GB=1000*1000*1000, G=1024*1024*1024, and so on for T, P, E, Z, Y.
 
-def parse_nbytes(nbytes):
+def parse_bytespec(nbytes):
+	smap = (('c',   1),
+			('w',   2),
+			('b',   512),
+			('kB',  1000),
+			('K',   1024),
+			('MB',  1000*1000),
+			('M',   1024*1024),
+			('GB',  1000*1000*1000),
+			('G',   1024*1024*1024),
+			('TB',  1000*1000*1000*1000),
+			('T',   1024*1024*1024*1024))
 	import re
-	m = re.match(r'([0123456789]+)(.*)',nbytes)
-	smap = ('c',1),('w',2),('b',512),('kB',1000),('K',1024),('MB',1000*1000),\
-			('M',1024*1024),('GB',1000*1000*1000),('G',1024*1024*1024)
+	m = re.match(r'([0123456789.]+)(.*)',nbytes)
 	if m:
 		if m.group(2):
 			for k,v in smap:
 				if k == m.group(2):
-					return int(m.group(1)) * v
+					from decimal import Decimal
+					return int(Decimal(m.group(1)) * v)
 			else:
 				msg("Valid byte specifiers: '{}'".format("' '".join([i[0] for i in smap])))
+		elif '.' in nbytes:
+			raise ValueError('fractional bytes not allowed')
 		else:
 			return int(nbytes)