From 7cc69a2725c65c80cd403aedd33b2522483c3ab7 Mon Sep 17 00:00:00 2001
From: MMGen <mmgen@tuta.io>
Date: Thu, 21 Mar 2019 11:46:50 +0000
Subject: [PATCH] mmgen-tool: rewrite 'rand2file', add test

---
 mmgen/seed.py             |  2 +-
 mmgen/tool.py             | 67 ++++++++++++++++++++-------------------
 test/test_py_d/ts_misc.py | 17 ++++++++++
 3 files changed, 52 insertions(+), 34 deletions(-)

diff --git a/mmgen/seed.py b/mmgen/seed.py
index 3ce39911..57d668b1 100755
--- a/mmgen/seed.py
+++ b/mmgen/seed.py
@@ -1042,7 +1042,7 @@ harder to find, you're advised to choose a much larger file size than this.
 					msg('File size must be an integer no less than {}'.format(min_fsize))
 
 				from mmgen.tool import MMGenToolCmd
-				MMGenToolCmd().rand2file(fn,str(fsize)) # threaded routine TODO: check safe
+				MMGenToolCmd().rand2file(fn,str(fsize))
 				check_offset = False
 			else:
 				die(1,'Exiting at user request')
diff --git a/mmgen/tool.py b/mmgen/tool.py
index eb52e10d..c91eb3b4 100755
--- a/mmgen/tool.py
+++ b/mmgen/tool.py
@@ -587,61 +587,62 @@ class MMGenToolCmdFileUtil(MMGenToolCmdBase):
 
 	def rand2file(self,outfile:str,nbytes:str,threads=4,silent=False):
 		"write 'n' bytes of random data to specified file"
-		nbytes = parse_bytespec(nbytes)
-		from Crypto import Random
-		rh = Random.new()
-		from queue import Queue
 		from threading import Thread
-		bsize = 2**20
-		roll = bsize * 4
-		if opt.outdir: outfile = make_full_path(opt.outdir,outfile)
-		f = open(outfile,'wb')
-
-		from Crypto.Cipher import AES
-		from Crypto.Util import Counter
-
-		key = get_random(32)
+		from queue import Queue
+		from cryptography.hazmat.primitives.ciphers import Cipher,algorithms,modes
+		from cryptography.hazmat.backends import default_backend
 
 		def encrypt_worker(wid):
+			ctr_init_val = os.urandom(g.aesctr_iv_len)
+			c = Cipher(algorithms.AES(key),modes.CTR(ctr_init_val),backend=default_backend())
+			encryptor = c.encryptor()
 			while True:
-				i,d = q1.get()
-				c = AES.new(key,AES.MODE_CTR,counter=Counter.new(g.aesctr_iv_len*8,initial_value=i))
-				enc_data = c.encrypt(d)
-				q2.put(enc_data)
+				q2.put(encryptor.update(q1.get()))
 				q1.task_done()
 
 		def output_worker():
 			while True:
-				data = q2.get()
-				f.write(data)
+				f.write(q2.get())
 				q2.task_done()
 
-		q1 = Queue()
+		nbytes = parse_bytespec(nbytes)
+		if opt.outdir:
+			outfile = make_full_path(opt.outdir,outfile)
+		f = open(outfile,'wb')
+
+		key = get_random(32)
+		q1,q2 = Queue(),Queue()
+
 		for i in range(max(1,threads-2)):
-			t = Thread(target=encrypt_worker,args=(i,))
+			t = Thread(target=encrypt_worker,args=[i])
 			t.daemon = True
 			t.start()
 
-		q2 = Queue()
 		t = Thread(target=output_worker)
 		t.daemon = True
 		t.start()
 
-		i = 1; rbytes = nbytes
-		while rbytes > 0:
-			d = rh.read(min(bsize,rbytes))
-			q1.put((i,d))
-			rbytes -= bsize
-			i += 1
-			if not (bsize*i) % roll:
-				msg_r('\rRead: {} bytes'.format(bsize*i))
+		blk_size = 1024 * 1024
+		for i in range(nbytes // blk_size):
+			if not i % 4:
+				msg_r('\rRead: {} bytes'.format(i * blk_size))
+			q1.put(os.urandom(blk_size))
+
+		if nbytes % blk_size:
+			q1.put(os.urandom(nbytes % blk_size))
 
-		if not silent:
-			msg('\rRead: {} bytes'.format(nbytes))
-			qmsg("\r{} bytes of random data written to file '{}'".format(nbytes,outfile))
 		q1.join()
 		q2.join()
 		f.close()
+
+		fsize = os.stat(outfile).st_size
+		if fsize != nbytes:
+			die(3,'{}: incorrect random file size (should be {})'.format(fsize,nbytes))
+
+		if not silent:
+			msg('\rRead: {} bytes'.format(nbytes))
+			qmsg("\r{} byte{} of random data written to file '{}'".format(nbytes,suf(nbytes),outfile))
+
 		return True
 
 class MMGenToolCmdWallet(MMGenToolCmdBase):
diff --git a/test/test_py_d/ts_misc.py b/test/test_py_d/ts_misc.py
index e462f98b..73a4c509 100755
--- a/test/test_py_d/ts_misc.py
+++ b/test/test_py_d/ts_misc.py
@@ -85,11 +85,28 @@ class TestSuiteTool(TestSuiteMain,TestSuiteBase):
 	enc_infn = 'tool_encrypt.in'
 	cmd_group = (
 		('tool_find_incog_data', (9,"'mmgen-tool find_incog_data'", [[[hincog_fn],1],[[incog_id_fn],1]])),
+		('tool_rand2file',       (9,"'mmgen-tool rand2file'", [])),
 		('tool_encrypt',         (9,"'mmgen-tool encrypt' (random data)",     [])),
 		('tool_decrypt',         (9,"'mmgen-tool decrypt' (random data)", [[[enc_infn+'.mmenc'],9]])),
 		# ('tool_encrypt_ref', (9,"'mmgen-tool encrypt' (reference text)",  [])),
 	)
 
+	def tool_rand2file(self):
+		outfile = os.path.join(self.tmpdir,'rand2file.out')
+		from mmgen.tool import MMGenToolCmd
+		tu = MMGenToolCmd()
+		for nbytes in ('1','1023','1K','1048575','1M','1048577','123M'):
+			t = self.spawn( 'mmgen-tool',
+							['-d',self.tmpdir,'-r0','rand2file','rand2file.out',nbytes],
+							extra_desc='({} byte{})'.format(nbytes,suf(tu.bytespec(nbytes)))
+							)
+			t.expect('random data written to file')
+			t.read()
+			t.p.wait()
+			t.ok()
+		t.skip_ok = True
+		return t
+
 	def tool_encrypt(self):
 		infile = joinpath(self.tmpdir,self.enc_infn)
 		write_to_file(infile,os.urandom(1033),binary=True)