import os
import tempfile
import struct
import re
from subprocess import Popen, PIPE
from nose.plugins.skip import Skip, SkipTest
import ubpf.assembler
import testdata
VM = os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "vm", "test")

def check_datafile(filename):
    """
    Given assembly source code and an expected result, run the eBPF program and
    verify that the result matches. Uses the JIT compiler.
    """
    data = testdata.read(filename)
    if 'asm' not in data and 'raw' not in data:
        raise SkipTest("no asm or raw section in datafile")
    if 'result' not in data and 'error' not in data and 'error pattern' not in data:
        raise SkipTest("no result or error section in datafile")
    if not os.path.exists(VM):
        raise SkipTest("VM not found")
    if 'no jit' in data:
        raise SkipTest("JIT disabled for this testcase (%s)" % data['no jit'])

    if 'raw' in data:
        code = ''.join(struct.pack("=Q", x) for x in data['raw'])
    else:
        code = ubpf.assembler.assemble(data['asm'])

    memfile = None

    cmd = [VM]
    if 'mem' in data:
        memfile = tempfile.NamedTemporaryFile()
        memfile.write(data['mem'])
        memfile.flush()
        cmd.extend(['-m', memfile.name])

    cmd.append('-j')
    cmd.append('-')

    vm = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)

    stdout, stderr = vm.communicate(code)
    stderr = stderr.strip()

    if memfile:
        memfile.close()

    if 'error' in data:
        if data['error'] != stderr:
            raise AssertionError("Expected error %r, got %r" % (data['error'], stderr))
    elif 'error pattern' in data:
        if not re.search(data['error pattern'], stderr):
            raise AssertionError("Expected error matching %r, got %r" % (data['error pattern'], stderr))
    else:
        if stderr:
            raise AssertionError("Unexpected error %r" % stderr)

    if 'result' in data:
        if vm.returncode != 0:
            raise AssertionError("VM exited with status %d, stderr=%r" % (vm.returncode, stderr))
        expected = int(data['result'], 0)
        result = int(stdout, 0)
        if expected != result:
            raise AssertionError("Expected result 0x%x, got 0x%x, stderr=%r" % (expected, result, stderr))
    else:
        if vm.returncode == 0:
            raise AssertionError("Expected VM to exit with an error code")

def test_datafiles():
    # Nose test generator
    # Creates a testcase for each datafile
    for filename in testdata.list_files():
        yield check_datafile, filename

@@ -1073,23 +1073,8 @@
 ne%0A%0A
-    cmd = %5BVM%5D%0A
     
@@ -1198,16 +1198,127 @@
 flush()%0A
+%0A    try:%0A        for register_offset in xrange(0, 10):%0A            cmd = %5BVM%5D%0A            if memfile:%0A        
         
@@ -1350,25 +1350,32 @@
 .name%5D)%0A
-%0A
+        
     cmd.
 append('
@@ -1370,42 +1370,65 @@
 cmd.
-app
+ext
 end(
+%5B
 '-j'
-)%0A    cmd.append(
+, '-r', str(register_offset), 
 '-'
+%5D
 )%0A%0A
+        
     
@@ -1482,16 +1482,24 @@
 =PIPE)%0A%0A
+        
     stdo
@@ -1532,16 +1532,24 @@
 e(code)%0A
+        
     stde
@@ -1577,45 +1577,12 @@
     
-if memfile:%0A        memfile.close()%0A%0A
+    
     
@@ -1593,32 +1593,40 @@
 error' in data:%0A
+        
         if data%5B
@@ -1637,32 +1637,40 @@
 or'%5D != stderr:%0A
+        
             rais
@@ -1737,24 +1737,32 @@
 %5D, stderr))%0A
+        
     elif 'er
@@ -1787,24 +1787,32 @@
 ta:%0A        
+       
+ 
 if not re.se
@@ -1848,16 +1848,24 @@
 tderr):%0A
+        
         
@@ -1965,30 +1965,44 @@
 r))%0A    
-else:%0A
+        else:%0A      
         
 if stder
@@ -1993,16 +1993,18 @@
         
+  
 if stder
@@ -1998,32 +1998,40 @@
      if stderr:%0A
+        
             rais
@@ -2080,16 +2080,24 @@
 tderr)%0A%0A
+        
     if '
@@ -2113,32 +2113,40 @@
 n data:%0A        
+      
+  
 if vm.returncode
@@ -2156,32 +2156,40 @@
  0:%0A            
+     
+   
 raise AssertionE
@@ -2250,32 +2250,40 @@
 ncode, stderr))%0A
+        
         expected
@@ -2316,16 +2316,24 @@
         
+  
+      
 result =
@@ -2356,16 +2356,24 @@
         
+     
+   
 if expec
@@ -2387,16 +2387,24 @@
 result:%0A
+        
         
@@ -2502,30 +2502,46 @@
 tderr))%0A    
-else:%0A
+        else:%0A        
         if v
@@ -2555,24 +2555,32 @@
 ncode == 0:%0A
+        
             
@@ -2641,16 +2641,77 @@
 r code%22)
+%0A    finally:%0A        if memfile:%0A            memfile.close()
 %0A%0Adef te

test JIT with each register offset