Repository: ntoll/uflash

Issue:
Add ability to extract Python code from a .hex file.
Because sometimes, you don't save the source file... ;-)

Discussion:
ntoll: I'll try to get this reviewed and merged this evening.
funkyHat: I've pushed another branch which fixes the test coverage (although one of the new tests is a bit of a beast...)
Also removed `uflash help` as mentioned above: https://github.com/funkyHat/uflash/tree/unhexlify-plus
ntoll: Can you merge the new branch into this one and I can do a final review..?
funkyHat: Done. Want me to squash it into a single commit?
ntoll: Please do... :-)

Patch:
diff --git a/uflash.py b/uflash.py
index 61fe83d..f5cd8a6 100644
--- a/uflash.py
+++ b/uflash.py
@@ -3,6 +3,7 @@
 This module contains functions for turning a Python script into a .hex file
 and flashing it onto a BBC micro:bit.
 """
+import argparse
 import sys
 import os
 import struct
@@ -19,8 +20,6 @@ _SCRIPT_ADDR = 0x3e000
 _HELP_TEXT = """
 Flash Python onto the BBC micro:bit
 
-Usage: uflash [path_to_script.py] [path_to_microbit]
-
 If no path to the micro:bit is provided uflash will attempt to autodetect the
 correct path to the device. If no path to the Python script is provided uflash
 will flash the unmodified MicroPython firmware onto the device.
@@ -72,6 +71,23 @@ def hexlify(script):
     return '\n'.join(output)
 
 
+def unhexlify(blob):
+    """
+    Takes a hexlified script and turns it back into Python code.
+    """
+    lines = blob.split('\n')[1:]
+    output = []
+    for line in lines:
+        # Discard the address, length etc. and reverse the hexlification
+        output.append(binascii.unhexlify(line[9:-2]))
+    # Strip off "MP<size>" from the start
+    output[0] = output[0][4:]
+    # and strip any null bytes from the end
+    output[-1] = output[-1].strip(b'\x00')
+    script = b''.join(output)
+    return script
+
+
 def embed_hex(runtime_hex, python_hex=None):
     """
     Given a string representing the MicroPython runtime hex, will embed a
@@ -98,6 +114,28 @@ def embed_hex(runtime_hex, python_hex=None):
     return '\n'.join(embedded_list) + '\n'
 
 
+def extract_script(embedded_hex):
+    """
+    Given a hex file containing the MicroPython runtime and an embedded Python
+    script, will extract the original script.
+
+    Returns a string containing the original embedded script.
+    """
+    hex_lines = embedded_hex.split('\n')
+    # Find the marker in the hex that comes just before the script
+    try:
+        start_line = hex_lines.index(':08058000193901005D150000AE') + 1
+    except ValueError as e:
+        raise ValueError('Bad input hex file:', e)
+    # Recombine the lines after that, but leave out the last 3 lines
+    blob = '\n'.join(hex_lines[start_line:-3])
+    if blob == '':
+        # If the result is the empty string, there was no embedded script
+        return b''
+    # Pass the extracted hex through unhexlify
+    return unhexlify(blob)
+
+
 def find_microbit():
     """
     Returns a path on the filesystem that represents the plugged in BBC
@@ -179,6 +217,8 @@ def flash(path_to_python=None, path_to_microbit=None):
     # Grab the Python script (if needed).
     python_hex = ''
     if path_to_python:
+        if not path_to_python.endswith('.py'):
+            raise ValueError('Python files must end in ".py".')
         with open(path_to_python, 'rb') as python_script:
             python_hex = hexlify(python_script.read())
     # Generate the resulting hex file.
@@ -195,6 +235,20 @@ def flash(path_to_python=None, path_to_microbit=None):
         raise IOError('Unable to find micro:bit. Is it plugged in?')
 
 
+def extract(path_to_hex=None, output_path=None):
+    """
+    Given a hex file this function will attempt to extract the embedded script
+    from it and save it either to output_path or stdout
+    """
+    with open(path_to_hex, 'r') as hex_file:
+        python_script = extract_script(hex_file.read())
+        if output_path is not None:
+            with open(output_path, 'w') as output_file:
+                output_file.write(python_script)
+        else:
+            print(python_script.decode('utf-8'))
+
+
 def main(argv=None):
     """
     Entry point for the command line tool 'uflash'.
@@ -210,20 +264,21 @@ def main(argv=None):
     """
     if not argv:
         argv = sys.argv[1:]
-    arg_len = len(argv)
     try:
-        if arg_len == 0:
-            flash()
-        elif arg_len >= 1:
-            if argv[0] == 'help':
-                print(_HELP_TEXT)
-                return
-            if not argv[0].lower().endswith('.py'):
-                raise ValueError('Python files must end in ".py".')
-            if arg_len == 1:
-                flash(argv[0])
-            elif arg_len > 1:
-                flash(argv[0], argv[1])
+        parser = argparse.ArgumentParser(description=_HELP_TEXT)
+        parser.add_argument('source', nargs='?', default=None)
+        parser.add_argument('target', nargs='?', default=None)
+        parser.add_argument('-e', '--extract',
+                            action='store_true',
+                            help="""Extract python source from a hex file
+                            instead of creating the hex file""",
+                            )
+        args = parser.parse_args(argv)
+
+        if args.extract:
+            extract(args.source, args.target)
+        else:
+            flash(args.source, args.target)
     except Exception as ex:
         # The exception of no return. Print the exception information.
         print(ex)


Test patch:
diff --git a/tests/test_uflash.py b/tests/test_uflash.py
index 5434084..eb0c87c 100644
--- a/tests/test_uflash.py
+++ b/tests/test_uflash.py
@@ -39,6 +39,15 @@ def test_hexlify():
     assert len(lines) == 5
 
 
+def test_unhexlify():
+    """
+    Ensure that we can get the script back out using unhexlify
+    """
+    hexlified = uflash.hexlify(TEST_SCRIPT)
+    unhexlified = uflash.unhexlify(hexlified)
+    assert unhexlified == TEST_SCRIPT
+
+
 def test_hexlify_empty_script():
     """
     The function returns an empty string if the script is empty.
@@ -84,6 +93,32 @@ def test_embed_no_runtime():
     assert ex.value.args[0] == 'MicroPython runtime hex required.'
 
 
+def test_extract():
+    """
+    The script should be returned if there is one
+    """
+    python = uflash.hexlify(TEST_SCRIPT)
+    result = uflash.embed_hex(uflash._RUNTIME, python)
+    extracted = uflash.extract_script(result)
+    assert extracted == TEST_SCRIPT
+
+
+def test_extract_not_valid_hex():
+    """
+    Return a sensible message if the hex file isn't valid
+    """
+    with pytest.raises(ValueError) as e:
+        uflash.extract_script('invalid input')
+    assert 'Bad input hex file' in e.value.args[0]
+
+
+def test_extract_no_python():
+    """
+    What to do here?
+    """
+    assert uflash.extract_script(uflash._RUNTIME) == b''
+
+
 def test_find_microbit_posix_exists():
     """
     Simulate being on os.name == 'posix' and a call to "mount" returns a
@@ -278,8 +313,7 @@ def test_main_no_args():
     with mock.patch('sys.argv', ['uflash', ]):
         with mock.patch('uflash.flash') as mock_flash:
             uflash.main()
-            assert mock_flash.call_count == 1
-            assert mock_flash.call_args == ()
+            assert mock_flash.called_once_with(None, None)
 
 
 def test_main_first_arg_python():
@@ -322,11 +356,55 @@ def test_main_two_args():
         assert mock_flash.called_once_with('foo.py', '/media/foo/bar')
 
 
-def test_main_extra_args_ignored():
+def test_extract_command():
     """
-    Any arguments more than two are ignored, with only the first two passed
-    into the flash() function.
+    Test the command-line script extract feature
     """
-    with mock.patch('uflash.flash') as mock_flash:
-        uflash.main(argv=['foo.py', '/media/foo/bar', 'baz', 'quux'])
-        assert mock_flash.called_once_with('foo.py', '/media/foo/bar')
+    with mock.patch('uflash.extract') as mock_extract:
+        uflash.main(argv=['-e', 'hex.hex', 'foo.py'])
+        assert mock_extract.called_once_with('hex.hex', 'foo.py')
+
+
+def test_extract_paths():
+    """
+    Test the different paths of the extract() function.
+    It should open and extract the contents of the file (input arg)
+    When called with only an input it should print the output of extract_script
+    When called with two arguments it should write the output to the output arg
+    """
+    mock_e = mock.MagicMock(return_value=mock.sentinel.script)
+    mock_o = mock.MagicMock()
+    mock_o.return_value.__enter__ = lambda s: s
+    mock_o.return_value.__exit__ = mock.Mock()
+    mock_o.return_value.read.return_value = 'script'
+    mock_o.return_value.write = mock.Mock()
+
+    with mock.patch('uflash.extract_script', mock_e) as mock_extract_script, \
+            mock.patch('builtins.print') as mock_print, \
+            mock.patch('builtins.open', mock_o) as mock_open:
+        uflash.extract('foo.hex')
+        assert mock_open.called_once_with('foo.hex')
+        assert mock_extract_script.called_once_with(mock.sentinel.file_handle)
+        assert mock_print.called_once_with(mock.sentinel.script)
+
+        uflash.extract('foo.hex', 'out.py')
+        assert mock_open.call_count == 3
+        assert mock_open.called_with('out.py', 'w')
+        assert mock_open.return_value.write.call_count == 1
+
+
+def test_extract_command_source_only():
+    """
+    If there is no target file the extract command should write to stdout
+    """
+    with mock.patch('uflash.extract') as mock_extract:
+        uflash.main(argv=['hex.hex'])
+        assert mock_extract.called_once_with('hex.hex')
+
+
+def test_extract_command_no_source():
+    """
+    If there is no source file the extract command should complain
+    """
+    with pytest.raises(TypeError):
+        uflash.extract(None, None)