diff --git a/mmgen/main.py b/mmgen/main.py index 4365d955..e2713cc2 100755 --- a/mmgen/main.py +++ b/mmgen/main.py @@ -27,7 +27,10 @@ def launch(what): if what == "keygen": what = "addrgen" try: import termios - except: __import__("mmgen.main_" + what) # Windows + except: # Windows + from mmgen.util import start_mscolor + start_mscolor() + __import__("mmgen.main_" + what) else: import sys,atexit fd = sys.stdin.fileno() diff --git a/mmgen/test.py b/mmgen/test.py index 60b302d0..c960b65c 100755 --- a/mmgen/test.py +++ b/mmgen/test.py @@ -22,17 +22,9 @@ test.py: Shared routines for the test suites import sys,os from binascii import hexlify -from mmgen.util import msg,write_to_file +from mmgen.util import msg,write_to_file,red,green import mmgen.opt as opt -_red,_grn,_yel,_cya,_reset = ( - ["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"] -) -def red(s): return _red+s+_reset -def green(s): return _grn+s+_reset -def yellow(s): return _yel+s+_reset -def cyan(s): return _cya+s+_reset - def cleandir(d): try: files = os.listdir(d) except: return @@ -57,8 +49,8 @@ def mk_tmpdir(cfg): def get_tmpfile_fn(cfg,fn): return os.path.join(cfg['tmpdir'],fn) -def write_to_tmpfile(cfg,fn,data): - write_to_file(os.path.join(cfg['tmpdir'],fn),data,silent=True) +def write_to_tmpfile(cfg,fn,data,mode='wb'): + write_to_file(os.path.join(cfg['tmpdir'],fn),data,silent=True,mode=mode) def read_from_tmpfile(cfg,fn): from mmgen.util import get_data_from_file diff --git a/mmgen/tool.py b/mmgen/tool.py index 4be2db2c..7a16a703 100755 --- a/mmgen/tool.py +++ b/mmgen/tool.py @@ -256,8 +256,11 @@ def hexdump(infile, cols=8, line_nums=True): cols=cols, line_nums=line_nums)) def unhexdump(infile): + if sys.platform[:3] == "win": + import msvcrt + msvcrt.setmode(sys.stdout.fileno(),os.O_BINARY) sys.stdout.write(decode_pretty_hexdump( - get_data_from_file(infile,dash=True,silent=True))) + get_data_from_file(infile,dash=True,silent=True))) def strtob58(s): enc = bitcoin.b58encode(s) @@ -566,7 +569,7 @@ def rand2file(outfile, nbytes, threads=4, silent=False): bsize = 2**20 roll = bsize * 4 if opt.outdir: outfile = make_full_path(opt.outdir,outfile) - f = open(outfile,"w") + f = open(outfile,"wb") from Crypto.Cipher import AES from Crypto.Util import Counter diff --git a/mmgen/util.py b/mmgen/util.py index 2dd78a8f..6ef081c1 100755 --- a/mmgen/util.py +++ b/mmgen/util.py @@ -29,15 +29,27 @@ import mmgen.globalvars as g pnm = g.proj_name -_red,_grn,_yel,_cya,_reset = ( - ["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"], - ["","","","",""] -)[int(sys.platform[:3] == "win")] +_red,_grn,_yel,_cya,_reset = \ + ["\033[%sm" % c for c in "31;1","32;1","33;1","36;1","0"] -def red(s): return _red+s+_reset -def green(s): return _grn+s+_reset -def yellow(s): return _yel+s+_reset -def cyan(s): return _cya+s+_reset +def red(s): return _red+s+_reset +def green(s): return _grn+s+_reset +def yellow(s): return _yel+s+_reset +def cyan(s): return _cya+s+_reset +def nocolor(s): return s + +def start_mscolor(): + if sys.platform[:3] == "win": + global red,green,yellow,cyan,nocolor + import os + if "MMGEN_NOMSCOLOR" in os.environ: + red = green = yellow = cyan = nocolor + else: + try: + import colorama + colorama.init(strip=True,convert=True) + except: + red = green = yellow = cyan = nocolor def msg(s): sys.stderr.write(s+"\n") def msg_r(s): sys.stderr.write(s) @@ -484,7 +496,8 @@ def write_to_file( desc="data", confirm_overwrite=False, verbose=False, - silent=False + silent=False, + mode='wb' ): if opt.outdir: outfile = make_full_path(opt.outdir,outfile) @@ -498,7 +511,7 @@ def write_to_file( else: if not silent: msg("Overwriting file '%s'" % outfile) - f = open_file_or_exit(outfile,'wb') + f = open_file_or_exit(outfile,mode) try: f.write(data) except: diff --git a/test/gentest.py b/test/gentest.py index 182547fb..84f9bf78 100755 --- a/test/gentest.py +++ b/test/gentest.py @@ -12,9 +12,11 @@ from binascii import hexlify import mmgen.opt as opt import mmgen.globalvars as g -from mmgen.util import msg,msg_r,mmsg,mdie,red,green,vmsg +from mmgen.util import msg,msg_r,mmsg,mdie,red,green,vmsg,start_mscolor from mmgen.bitcoin import hextowif,privnum2addr +start_mscolor() + rounds = 100 opts_data = { 'desc': "Test addresses generated by {} against output of 'keyconv'".format(g.proj_name), diff --git a/test/test.py b/test/test.py index 6fe23116..573921c7 100755 --- a/test/test.py +++ b/test/test.py @@ -17,16 +17,11 @@ scripts = ( import mmgen.globalvars as g import mmgen.opt as opt -from mmgen.util import mmsg,mdie,Msg,die,capfirst,write_data_to_file -from mmgen.test import * +from mmgen.util import * -if sys.platform[:3] == "win": - try: - import colorama - colorama.init(strip=True,convert=True) - except: - def nocolor(s): return s - red = green = yellow = cyan = nocolor +start_mscolor() + +from mmgen.test import * tb_cmd = "scripts/traceback.py" hincog_fn = "rand_data" @@ -420,8 +415,8 @@ cmd_args = opt.opts.init(opts_data) if opt.system: sys.path.pop(0) ni = bool(opt.non_interactive) -# temporary -#os.environ["MMGEN_USE_OLD_SCRIPTS"] = "1" +# Disable MS color in spawned scripts due to bad interactions +os.environ["MMGEN_NOMSCOLOR"] = "1" if opt.debug_scripts: os.environ["MMGEN_DEBUG"] = "1" diff --git a/test/tooltest.py b/test/tooltest.py index f3831104..c6aa6f55 100755 --- a/test/tooltest.py +++ b/test/tooltest.py @@ -9,9 +9,11 @@ os.chdir(os.path.join(pn,os.pardir)) sys.path.__setitem__(0,os.path.abspath(os.curdir)) import mmgen.opt as opt -from mmgen.util import msg,msg_r,vmsg,vmsg_r,Msg,mmsg,mdie +from mmgen.util import msg,msg_r,vmsg,vmsg_r,Msg,mmsg,mdie,start_mscolor from collections import OrderedDict +start_mscolor() + cmd_data = OrderedDict([ ('util', { 'desc': "base conversion, hashing and file utilities", @@ -150,12 +152,12 @@ class MMGenToolTestSuite(object): self.__class__.__dict__[cmd](*([self,cmd] + file_list)) - def run_cmd(self,name,tool_args,kwargs="",extra_msg="",silent=False): + def run_cmd(self,name,tool_args,kwargs="",extra_msg="",silent=False,strip=True): mmgen_tool = "mmgen-tool" if not opt.system: mmgen_tool = os.path.join(os.curdir,mmgen_tool) - sys_cmd = [mmgen_tool, "-d",cfg['tmpdir'], name] + tool_args + kwargs.split() + sys_cmd = ["python", mmgen_tool, "-d",cfg['tmpdir'], name] + tool_args + kwargs.split() if extra_msg: extra_msg = "(%s)" % extra_msg full_name = " ".join([name]+kwargs.split()+extra_msg.split()) if not silent: @@ -166,12 +168,13 @@ class MMGenToolTestSuite(object): msg_r("Testing %-31s%s" % (full_name+":","")) import subprocess - return subprocess.check_output(sys_cmd) + ret = subprocess.check_output(sys_cmd) + return (ret,ret.rstrip())[int(strip)] def run_cmd_chk(self,name,f1,f2,kwargs="",extra_msg=""): - idata = read_from_file(f1)[:-1] - odata = read_from_file(f2)[:-1] - ret = self.run_cmd(name,[odata],kwargs=kwargs,extra_msg=extra_msg)[:-1] + idata = read_from_file(f1).rstrip() + odata = read_from_file(f2).rstrip() + ret = self.run_cmd(name,[odata],kwargs=kwargs,extra_msg=extra_msg) vmsg("In: " + repr(odata)) vmsg("Out: " + repr(ret)) if ret == idata: ok() @@ -182,32 +185,32 @@ class MMGenToolTestSuite(object): return ret def run_cmd_nochk(self,name,f1,kwargs=""): - odata = read_from_file(f1)[:-1] - ret = self.run_cmd(name,[odata],kwargs=kwargs)[:-1] + odata = read_from_file(f1).rstrip() + ret = self.run_cmd(name,[odata],kwargs=kwargs) vmsg("In: " + repr(odata)) vmsg("Out: " + repr(ret)) return ret def run_cmd_out(self,name,carg=None,Return=False,kwargs="",fn_idx="",extra_msg=""): - if carg: write_to_tmpfile(cfg,"%s%s.in" % (name,fn_idx),carg+"\n") + if carg: write_to_tmpfile(cfg,"%s%s.in" % (name,fn_idx),carg+"\n",mode='w') ret = self.run_cmd(name,[carg] if carg else [],kwargs=kwargs,extra_msg=extra_msg) if carg: vmsg("In: " + repr(carg)) - vmsg("Out: " + repr(ret[:-1])) + vmsg("Out: " + repr(ret)) if ret: - write_to_tmpfile(cfg,"%s%s.out" % (name,fn_idx),ret) + write_to_tmpfile(cfg,"%s%s.out" % (name,fn_idx),ret+"\n",mode='w') if Return: return ret else: ok() else: msg(red("Error for command '%s'" % name)) sys.exit(3) - def run_cmd_randfileinput(self,name): + def run_cmd_randinput(self,name,strip=True): s = os.urandom(128) fn = name+".in" write_to_tmpfile(cfg,fn,s) - ret = self.run_cmd(name,[get_tmpfile_fn(cfg,fn)]) + ret = self.run_cmd(name,[get_tmpfile_fn(cfg,fn)],strip=strip) fn = name+".out" - write_to_tmpfile(cfg,fn,ret) + write_to_tmpfile(cfg,fn,ret+"\n",mode='w') ok() vmsg("Returned: %s" % ret) @@ -235,19 +238,19 @@ class MMGenToolTestSuite(object): self.run_cmd(name,[]) ok() - def id6(self,name): self.run_cmd_randfileinput(name) - def id8(self,name): self.run_cmd_randfileinput(name) - def hexdump(self,name): self.run_cmd_randfileinput(name) + def id6(self,name): self.run_cmd_randinput(name) + def id8(self,name): self.run_cmd_randinput(name) + def hexdump(self,name): self.run_cmd_randinput(name,strip=False) def unhexdump(self,name,fn1,fn2): - ret = self.run_cmd(name,[fn2]) + ret = self.run_cmd(name,[fn2],strip=False) orig = read_from_file(fn1) cmp_or_die(orig,ret) def rand2file(self,name): of = name + ".out" dlen = 1024 - self.run_cmd(name,[of,str(1024),"threads=4","silent=1"]) + self.run_cmd(name,[of,str(1024),"threads=4","silent=1"],strip=False) d = read_from_tmpfile(cfg,of) cmp_or_die(dlen,len(d)) @@ -262,14 +265,14 @@ class MMGenToolTestSuite(object): def b32tohex(self,name,f1,f2): self.run_cmd_chk(name,f1,f2) def b58randenc(self,name): ret = self.run_cmd_out(name,Return=True) - ok_or_die(ret[:-1],is_b58_str,"base 58 string") + ok_or_die(ret,is_b58_str,"base 58 string") def randhex(self,name): ret = self.run_cmd_out(name,Return=True) - ok_or_die(ret[:-1],binascii.unhexlify,"hex string") + ok_or_die(ret,binascii.unhexlify,"hex string") def randwif(self,name): for n,k in enumerate(["","compressed=1"]): ret = self.run_cmd_out(name,kwargs=k,Return=True,fn_idx=n+1) - ok_or_die(ret[:-1],is_wif,"WIF key") + ok_or_die(ret,is_wif,"WIF key") def randpair(self,name): for n,k in enumerate(["","compressed=1"]): wif,addr = self.run_cmd_out(name,kwargs=k,Return=True,fn_idx=n+1).split()