Examples

carmeli-tamir · Nov 24, 2019 · c4aabd0 · c4aabd0
1 parent eddebc1
commit c4aabd0
Show file tree

Hide file tree

Showing 4 changed files with 84 additions and 76 deletions.
diff --git a/TODO b/TODO
@@ -1,7 +1,7 @@
-- Make travis ci
-    - Implement tests for all 3 cases
-    - combine tests and setup?
 - Improve documentation and open project to the public
+    - Readme structure
+    - Readme how to set up
+    - list of issues
 - Reduce permissions for /dev/skunk
 - Re-write skunk.py as a class
 - Generate parts of skunk.h automatically

diff --git a/examples/examples1.py b/examples/examples1.py
@@ -0,0 +1,25 @@
+from skunk import Skunk
+import skunk_pb2
+
+def run_skunk():
+    sk = Skunk()
+
+    ret = sk.call_function_two_arg("kallsyms_lookup_name",
+        1, skunk_pb2.FunctionCall.eight_byte,
+        "kallsyms_lookup_name", skunk_pb2.Argument.string)
+    print("The address of kallsyms_lookup_name is {}".format(hex(2**64 + ret.eight_byte))) # Printing 2's complement of the address
+
+    ret = sk.call_function_two_arg("round_jiffies",
+        1, skunk_pb2.FunctionCall.eight_byte,
+        133713371337, skunk_pb2.Argument.eight_byte)
+
+    print("The result of round_jiffies(133713371337) is {}".format(ret.eight_byte))
+
+    ret = sk.call_function_two_arg("strstr",
+        2, skunk_pb2.FunctionCall.string,
+        "whatisthemeaningoflife", skunk_pb2.Argument.string, "ning", skunk_pb2.Argument.string)
+
+    print("strstr(\"whatisthemeaningoflife\", \"ning\") == \"{}\"".format(ret.string))
+
+if __name__ == "__main__":
+    run_skunk()
diff --git a/user/skunk.py b/user/skunk.py
@@ -4,63 +4,44 @@
 
 import skunk_pb2
 
-def binary_length_and_value(buffer, buffer_length):
-    return struct.pack('I', buffer_length) + buffer
-
-def call_function_two_arg(device, fname, num_args, fret, arg1, type1, arg2=None, type2=None):
-    func_with_one_arg = skunk_pb2.FunctionCall()
-    func_with_one_arg.returnType = fret
-    func_with_one_arg.numberOfArguments = num_args
-    func_with_one_arg.name = fname
-    func_with_one_arg.arg1.type = type1
-
-    if type1 == skunk_pb2.Argument.string:
-        func_with_one_arg.arg1.arg_string = arg1
-    elif type1 == skunk_pb2.Argument.eight_byte:
-        func_with_one_arg.arg1.arg_eight_byte = arg1
-    else:
-        raise ValueError("Unsupported argument type")
-
-    if type2 is not None:
-        func_with_one_arg.arg2.type = type2
-        if type2 == skunk_pb2.Argument.string:
-            func_with_one_arg.arg2.arg_string = arg2
-        elif type2 == skunk_pb2.Argument.eight_byte:
-            func_with_one_arg.arg2.arg_eight_byte = arg2
+class Skunk(object):
+    def __init__(self, device_name="/dev/skunk", ioctl_num=0xc008ee00):
+        self.device_name = device_name
+        self.ioctl_num = ioctl_num
+
+    def call_function_two_arg(self, fname, num_args, fret, arg1, type1, arg2=None, type2=None):
+        func_with_one_arg = skunk_pb2.FunctionCall()
+        func_with_one_arg.returnType = fret
+        func_with_one_arg.numberOfArguments = num_args
+        func_with_one_arg.name = fname
+        self._assign_arg(func_with_one_arg.arg1, type1, arg1)
+
+        if arg2 is not None and type2 is not None:
+            self._assign_arg(func_with_one_arg.arg2, type2, arg2)
+
+        func_with_one_arg_binary = self._binary_length_and_value(func_with_one_arg.SerializeToString(), func_with_one_arg.ByteSize())
+
+        with open(self.device_name, 'r') as skunk_device:
+            call_result = fcntl.ioctl(skunk_device, self.ioctl_num, func_with_one_arg_binary)
+
+        length_of_size_field = struct.calcsize('I')
+        return_length, = struct.unpack('I', call_result[:length_of_size_field])
+
+        return_value = skunk_pb2.ReturnValue()
+        return_value.ParseFromString(call_result[length_of_size_field:length_of_size_field + return_length])
+
+        return return_value
+
+    @staticmethod
+    def _assign_arg(protobuf_arg, arg_type, arg_value):
+        protobuf_arg.type = arg_type
+        if arg_type == skunk_pb2.Argument.string:
+            protobuf_arg.arg_string = arg_value
+        elif arg_type == skunk_pb2.Argument.eight_byte:
+            protobuf_arg.arg_eight_byte = arg_value
         else:
             raise ValueError("Unsupported argument type")
-
-    func_with_one_arg_binary = binary_length_and_value(func_with_one_arg.SerializeToString(), func_with_one_arg.ByteSize())
-
-    call_result = fcntl.ioctl(device, 0xc008ee00, func_with_one_arg_binary)
-
-    length_of_size_field = struct.calcsize('I')
-    return_length, = struct.unpack('I', call_result[:length_of_size_field])
-
-    return_value = skunk_pb2.ReturnValue()
-    return_value.ParseFromString(call_result[length_of_size_field:length_of_size_field + return_length])
-
-    return return_value
-
-
-def run_skunk():
-    with open("/dev/skunk", 'r') as skunk_device:
-        ret = call_function_two_arg(skunk_device, "kallsyms_lookup_name",
-        1, skunk_pb2.FunctionCall.eight_byte,
-        "kallsyms_lookup_name", skunk_pb2.Argument.string)
-        print("Got adress of {}".format(hex(2**64 + ret.eight_byte))) # Printing 2's complement of the address
-
-        ret = call_function_two_arg(skunk_device, "round_jiffies",
-        1, skunk_pb2.FunctionCall.eight_byte,
-        133713371337, skunk_pb2.Argument.eight_byte)
-
-        print("Got result of {}".format(ret.eight_byte))
-
-        ret = call_function_two_arg(skunk_device, "strstr",
-        2, skunk_pb2.FunctionCall.string,
-        "whatisthemeaningoflife", skunk_pb2.Argument.string, "ning", skunk_pb2.Argument.string)
-
-        print("Got result of {}".format(ret.string))
-
-if __name__ == "__main__":
-    run_skunk()
+
+    @staticmethod
+    def _binary_length_and_value(buffer, buffer_length):
+        return struct.pack('I', buffer_length) + buffer
diff --git a/user/test/test_basic.py b/user/test/test_basic.py
@@ -1,12 +1,17 @@
 import unittest
+from unittest import mock
 from unittest.mock import patch
 
-import skunk
+from skunk import Skunk
 import skunk_pb2
 
 
 class TestCalls(unittest.TestCase):
 
+    @classmethod
+    def setup_class(cls):
+        cls.skunk = Skunk("/dev/null")
+
     @staticmethod
     def build_return_value(status, eight_byte=None, string=None):
         ret = skunk_pb2.ReturnValue()
@@ -19,7 +24,7 @@ def build_return_value(status, eight_byte=None, string=None):
         else:
             ret.has_string = False
 
-        return skunk.binary_length_and_value(ret.SerializeToString(), ret.ByteSize())
+        return Skunk._binary_length_and_value(ret.SerializeToString(), ret.ByteSize())
 
     @staticmethod
     def assign_arg(protobuf_arg, arg_type, arg_value):
@@ -44,15 +49,14 @@ def build_function_call(return_type, function_name, *args):
         if len(args) > 2:
             TestCalls.assign_arg(func_with_one_arg.arg2, args[2], args[3])
 
-        return skunk.binary_length_and_value(func_with_one_arg.SerializeToString(), func_with_one_arg.ByteSize())
+        return Skunk._binary_length_and_value(func_with_one_arg.SerializeToString(), func_with_one_arg.ByteSize())
 
     @patch('skunk.fcntl')
     def test_call_kallsyms(self, mock_requests):
         ioctl_ret = self.build_return_value(skunk_pb2.ReturnValue.Success, 0x1ee7, None)
         mock_requests.ioctl.return_value = ioctl_ret
 
-        ret = skunk.call_function_two_arg(
-            None,
+        ret = TestCalls.skunk.call_function_two_arg(
             "kallsyms_lookup_name",
             1,
             skunk_pb2.FunctionCall.eight_byte,
@@ -61,8 +65,8 @@ def test_call_kallsyms(self, mock_requests):
         )
 
         mock_requests.ioctl.assert_called_once_with(
-                None,
-                0xc008ee00, 
+                mock.ANY,
+                mock.ANY, 
                 self.build_function_call(skunk_pb2.FunctionCall.eight_byte, "kallsyms_lookup_name", skunk_pb2.Argument.string,  "kallsyms_lookup_name")
             )
         assert ret.eight_byte == 0x1ee7
@@ -72,8 +76,7 @@ def test_call_roundjiffies(self, mock_requests):
         ioctl_ret = self.build_return_value(skunk_pb2.ReturnValue.Success, 133713371500, None)
         mock_requests.ioctl.return_value = ioctl_ret
 
-        ret = skunk.call_function_two_arg(
-            None,
+        ret = TestCalls.skunk.call_function_two_arg(
             "round_jiffies",
             1,
             skunk_pb2.FunctionCall.eight_byte,
@@ -82,8 +85,8 @@ def test_call_roundjiffies(self, mock_requests):
         )
 
         mock_requests.ioctl.assert_called_once_with(
-                None,
-                0xc008ee00, 
+                mock.ANY,
+                mock.ANY, 
                 self.build_function_call(skunk_pb2.FunctionCall.eight_byte, "round_jiffies", skunk_pb2.Argument.eight_byte,  133713371337)
             )
         assert ret.eight_byte == 133713371500
@@ -93,8 +96,7 @@ def test_call_strstr(self, mock_requests):
         ioctl_ret = self.build_return_value(skunk_pb2.ReturnValue.Success, None, "ningoflife")
         mock_requests.ioctl.return_value = ioctl_ret
 
-        ret = skunk.call_function_two_arg(
-            None,
+        ret = TestCalls.skunk.call_function_two_arg(
             "strstr",
             2,
             skunk_pb2.FunctionCall.string,
@@ -105,8 +107,8 @@ def test_call_strstr(self, mock_requests):
         )
 
         mock_requests.ioctl.assert_called_once_with(
-                None,
-                0xc008ee00, 
+                mock.ANY,
+                mock.ANY, 
                 self.build_function_call(skunk_pb2.FunctionCall.string, "strstr", skunk_pb2.Argument.string,  "whatisthemeaningoflife", skunk_pb2.Argument.string,  "ning")
             )
         assert ret.string == "ningoflife"