file: Change from class to functions

Fixes: #92

Signed-off-by: John Andersen <[email protected]>

Author: John Andersen

cve_bin_tool/cli.py

@@ -19,13 +19,12 @@
 import subprocess
 import logging
 import argparse
-import binascii
 import pkg_resources
 
 from .util import DirWalk, inpath
 from .extractor import Extractor
 from .strings import Strings
-from .file import File
+from .file import is_binary
 from .NVDAutoUpdate import NVDSQLite
 from .log import LOGGER
 
@@ -105,8 +104,7 @@ def scan_file(self, filename):
             raise InvalidFileError(filename)
 
         # step 1: check if it's an ELF binary file
-        f = File(filename)
-        if not f.check():
+        if not is_binary(filename):
             return None
         # parse binary file's strings
         s = Strings(filename)

cve_bin_tool/file.py

@@ -1,50 +1,31 @@
-#!/usr/bin/python3
-
 """
-This is the customized file class, by doing this
-the tool is able to support other operating systems like Windows
-and MacOS.
+To determine if a file is executable we read the first few bytes of a file and
+check it against various signatures for different executable file formats.
 """
-
-import string
 import sys
-import platform
+import inspect
 import binascii
 
-class File(object):
-
-    def __init__(self, filename = ''):
-        self.filename = filename
-        self.output = ''
-        self.os = platform.system()
-        self.signature = ''
-        if filename == '':
-            pass
-        else:
-            self.signature = self.read_signature().decode('utf-8')
-
-    def read_signature(self):
-        # read signature (first 4 bytes) of the file
-        with open(self.filename, 'rb') as f:
-            byte = f.read(4)
-            return binascii.hexlify(byte)
-
-    def check(self):
-        # make sure class has signature before check
-        if self.signature == '':
-            self.signature = self.read_signature().decode('utf-8')
-        return self.check_linux() or self.check_windows()
-
-    def check_linux(self):
-        # by decoding we could ignore the affect caused by python version
-        if sys.byteorder == 'little':
-            return self.signature == '7f454c46'
-        else:
-            return self.signature == '457f464c'
-
-    def check_windows(self):
-        # in windows/dos, the signature of exectuable is 0x5a4d
-        if sys.byteorder == 'little':
-            return self.signature[:4] == '4d5a'
-        else:
-            return self.signature[:4] == '5a4d'
+def is_binary(filename):
+    """ Read the magic bytes from a file and determine if it is an executable
+    binary."""
+    signature = read_signature(filename)
+    for name, method in inspect.getmembers(sys.modules[__name__],
+                                           predicate=inspect.isfunction):
+        if name.startswith('check_'):
+            if method(filename, signature):
+                return True
+    return False
+
+def read_signature(filename, length=4):
+    """ Read the signature, first length bytes, from filename."""
+    with open(filename, 'rb') as file_handle:
+        return file_handle.read(length)
+
+def check_elf(_filename, signature):
+    """ Check for an ELF signature."""
+    return signature == b'\x7f\x45\x4c\x46'
+
+def check_pe(_filename, signature):
+    """ Check for windows/dos PE signature, aka 0x5a4d."""
+    return signature[:4] == b'\x4d\x5a'

test/assests/single-byte.txt

@@ -0,0 +1 @@
+1

test/assests/windows.txt

@@ -0,0 +1 @@
+MZ

test/test_file.py

@@ -1,32 +1,32 @@
 """
-CVE-bin-tool File tests
+CVE-bin-tool file tests
 """
 import os
 import sys
 import unittest
 import subprocess
 
-from cve_bin_tool.file import File
+from cve_bin_tool.file import is_binary
 
+ASSETS_PATH = os.path.join(
+    os.path.abspath(os.path.dirname(__file__)), 'assests')
 BINARIES_PATH = os.path.join(
     os.path.abspath(os.path.dirname(__file__)), 'binaries')
 
-
 class TestFile(unittest.TestCase):
-    """ Tests the CVE Bin Tool Strings"""
+    """ Tests the CVE Bin Tool file binary checker."""
 
     @classmethod
     def setUpClass(cls):
         # build binaries
         subprocess.call(["make", "clean"], cwd=BINARIES_PATH)
         subprocess.call(["make", "all"], cwd=BINARIES_PATH)
-        cls.file = File()
 
     def _check_test(self, filename, is_executable):
         """Helper function to parse a binary file and check whether
         the given string is in the parsed result"""
-        self.file.filename = os.path.join(BINARIES_PATH, filename)
-        self.assertEqual(self.file.check(), is_executable)
+        self.assertEqual(is_binary(os.path.join(BINARIES_PATH, filename)),
+                         is_executable)
 
     def test_curl_7_34_0_out(self):
         """file test-curl-7.34.0.out"""
@@ -35,3 +35,12 @@ def test_curl_7_34_0_out(self):
     def test_curl_7_34_0_source(self):
         """file test-curl-7.34.0.c"""
         self._check_test('test-curl-7.34.0.c', False)
+
+    def test_single_byte_file(self):
+        """file single-byte"""
+        self.assertFalse(is_binary(os.path.join(ASSETS_PATH,
+                                                'single-byte.txt')))
+
+    def test_windows(self):
+        """file single-byte"""
+        self.assertTrue(is_binary(os.path.join(ASSETS_PATH, 'windows.txt')))