From d5242df5cfa98aaf77918fdfc7c75d8af77bf20a Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Fri, 27 Oct 2017 22:30:46 +0700
Subject: [PATCH] Add more tests

---
 myth                        | 21 +++++++++++----------
 tests/cmline_test.py        | 19 +++++++++++++++++++
 tests/testdata/metacoin.sol | 15 +++++++++++++++
 3 files changed, 45 insertions(+), 10 deletions(-)
 create mode 100644 tests/cmline_test.py
 create mode 100644 tests/testdata/metacoin.sol

diff --git a/myth b/myth
index a3e69ea4..b47350e3 100755
--- a/myth
+++ b/myth
@@ -13,6 +13,7 @@ from mythril.rpc.client import EthJsonRpc
 from mythril.ipc.client import EthIpc
 from ethereum import utils
 from laser.ethereum import laserfree
+from pathlib import Path
 import logging
 import sys
 import argparse
@@ -33,21 +34,22 @@ def exitWithError(message):
 
 
 parser = argparse.ArgumentParser(description='Bug hunting on the Ethereum blockchain')
+parser.add_argument("solidity_file", nargs='*')
+
 
 commands = parser.add_argument_group('commands')
 commands.add_argument('-d', '--disassemble',  action='store_true', help='disassemble')
 commands.add_argument('-g', '--graph', help='generate a control flow graph', metavar='OUTPUT_FILE')
 commands.add_argument('-x', '--fire-lasers', action='store_true', help='detect vulnerabilities')
-commands.add_argument('-t', '--trace', action='store_true', help='trace, use with -c or -a and --data (optional)')
-commands.add_argument('-s', '--search', help='search the contract database')
-commands.add_argument('--xrefs', action='store_true', help='get xrefs from a contract, use with -c or -a')
+commands.add_argument('-t', '--trace', action='store_true', help='trace contract, use with --data (optional)')
+commands.add_argument('-s', '--search', help='search the contract database', metavar='EXPRESSION')
+commands.add_argument('--xrefs', action='store_true', help='get xrefs from a contract')
 commands.add_argument('--hash', help='calculate function signature hash', metavar='SIGNATURE')
 commands.add_argument('--init-db', action='store_true', help='initialize the contract database')
 
 inputs = parser.add_argument_group('input arguments')
 inputs.add_argument('-c', '--code', help='hex-encoded bytecode string ("6060604052...")', metavar='BYTECODE')
-inputs.add_argument('-i', '--infiles', help='comma-separated list of solidity files', metavar='INPUT_FILES')
-inputs.add_argument('-a', '--address', help='pull contract from the mainnet', metavar='CONTRACT_ADDRESS')
+inputs.add_argument('-a', '--address', help='pull contract from the blockchain', metavar='CONTRACT_ADDRESS')
 inputs.add_argument('--data', help='message call input data for tracing')
 
 options = parser.add_argument_group('options')
@@ -105,7 +107,7 @@ if args.search or args.init_db:
 
 # Establish RPC/IPC connection if necessary
 
-if (args.address or args.infiles):
+if (args.address or len(args.solidity_file)):
         if args.ipc:
             try:
                 eth = EthIpc()
@@ -127,13 +129,12 @@ if (args.code):
     contracts.append(ETHContract(args.code, address = util.get_random_address()))
 elif (args.address):
     contracts.append(ETHContract(eth.eth_getCode(args.address), address = args.address))
-elif (args.infiles):
-    files = args.infiles.split(",")
-
+elif (len(args.solidity_file)):
     counter = 0
 
-    for file in files:
+    for file in args.solidity_file:
 
+        file = file.replace("~", str(Path.home())) # Expand user path
         name, bytecode = compile_solidity(solc_binary, file)
 
         # Max. 16 contracts supported
diff --git a/tests/cmline_test.py b/tests/cmline_test.py
new file mode 100644
index 00000000..93e044b5
--- /dev/null
+++ b/tests/cmline_test.py
@@ -0,0 +1,19 @@
+import unittest
+import os
+from subprocess import check_output
+
+
+class CommandLineToolTestCase(unittest.TestCase):
+
+    def runTest(self):
+
+        script_path = os.path.dirname(os.path.realpath(__file__))
+        myth = os.path.join(script_path, '..', 'myth')
+
+        out = check_output([myth,'-d','-c', '0x5050']).decode("UTF-8")
+
+        self.assertEqual('0 POP\n1 POP\n', out)
+
+        out = check_output([myth,'-d', os.path.join(script_path,'testdata','metacoin.sol')]).decode("UTF-8")
+
+        self.assertIn('0 PUSH1 0x60\n2 PUSH1 0x40', out)
\ No newline at end of file
diff --git a/tests/testdata/metacoin.sol b/tests/testdata/metacoin.sol
new file mode 100644
index 00000000..e5efbf4d
--- /dev/null
+++ b/tests/testdata/metacoin.sol
@@ -0,0 +1,15 @@
+pragma solidity ^0.4.17;
+
+contract metaCoin {
+	mapping (address => uint) public balances;
+	function metaCoin() public {
+		balances[msg.sender] = 10000;
+	}
+	
+	function sendToken(address receiver, uint amount) public returns(bool successful){
+		if (balances[msg.sender] < amount) return false;
+		balances[msg.sender] -= amount;
+		balances[receiver] += amount;
+		return false;
+	}
+}