Added support for downloading pre-compiled signatures

wordfence/api/noc1.py

@@ -1,13 +1,18 @@
 import json
 import re
+import base64
 from typing import Callable, Optional
 
 from .noc_client import NocClient
 from .exceptions import ApiException
 from .licensing import License
 
-from ..intel.signatures import CommonString, Signature, SignatureSet
-from ..util.validation import DictionaryValidator, ListValidator, Validator
+from ..intel.signatures import CommonString, Signature, SignatureSet, \
+    PrecompiledSignatureSet
+from ..util.validation import DictionaryValidator, ListValidator, Validator, \
+    OptionalValueValidator
+from ..util.platform import Platform
+from ..util.serialization import limited_deserialize
 
 NOC1_BASE_URL = 'https://noc1.wordfence.com/v2.27/'
 
@@ -137,6 +142,58 @@ def get_malware_signatures(self) -> SignatureSet:
                         ) from index_error
         return SignatureSet(common_strings, signatures, self.license)
 
+    def get_precompiled_patterns(
+                self,
+                platform: str,
+                library_version: str,
+                library_type: Optional[str] = None,
+                database_version: int = PrecompiledSignatureSet.VERSION
+            ) -> dict:
+        parameters = {
+                'platform': platform,
+                'library_version': library_version,
+                'database_version': database_version
+            }
+        if library_type is not None:
+            parameters['library_type'] = library_type
+        response = self.request('get_precompiled_patterns', parameters)
+        validator = DictionaryValidator({
+                'data': OptionalValueValidator(str)
+            })
+        self.validate_response(response, validator)
+        return response
+
+    def get_precompiled_malware_signatures(
+                self,
+                platform: Platform,
+                library_version: str,
+                library_type: Optional[str] = None,
+                database_version: int = PrecompiledSignatureSet.VERSION
+            ) -> Optional[PrecompiledSignatureSet]:
+        response = self.get_precompiled_patterns(
+                platform.key,
+                library_version,
+                library_type,
+                database_version
+            )
+        data = response['data']
+        if data is None:
+            return None
+        data = base64.b64decode(data)
+        signature_set = limited_deserialize(
+                response.data,
+                {
+                    'wordfence.intel.signatures.PrecompiledSignatureSet',
+                    'wordfence.intel.signatures.SignatureSet',
+                    'wordfence.intel.signatures.Signature'
+                }
+            )
+        if isinstance(signature_set, PrecompiledSignatureSet):
+            return signature_set
+        raise ApiException(
+                'Malformed signature set data received from Wordfence API'
+            )
+
     def ping_api_key(self) -> bool:
         return self.process_simple_request('ping_api_key')
 

wordfence/api/noc_client.py

@@ -21,6 +21,8 @@ def __init__(
         self.base_url = base_url \
             if base_url is not None \
             else self.get_default_base_url()
+        from wordfence.logging import log
+        log.debug(f'NOC1 Base URL: {self.base_url}')
         self.timeout = timeout
 
     def get_default_base_url(self) -> str:

wordfence/cli/context.py

@@ -181,8 +181,10 @@ def display_version(self) -> None:
         has_vectorscan = self.has_vectorscan()
         vectorscan_support_text = yes_no(has_vectorscan)
         if has_vectorscan:
-            vectorscan_support_text += \
-                f' - Version: {vectorscan.VERSION}'
+            vectorscan_support_text += (
+                    f' - Version: {vectorscan.VERSION} (API Version: '
+                    f'{vectorscan.API_VERSION})'
+                )
         print(f'Vectorscan Supported: {vectorscan_support_text}')
 
     def has_terminal_output(self) -> bool:

wordfence/cli/malwarescan/definition.py

@@ -204,14 +204,23 @@
         "argument_type": "FLAG",
         "default": False
     },
+    "pre-compile-generic": {
+        "description": "Pre-compile and cache the signature set without "
+                       "any CPU-specific optimizations and without running "
+                       "a scan",
+        "context": "CLI",
+        "argument_type": "FLAG",
+        "default": False
+    },
     "pattern-database-path": {
         "description": "Use an alternate path for storage of the pattern "
                        "database",
         "context": "ALL",
         "argument_type": "OPTION",
         "meta": {
             "accepts_file": True
-        }
+        },
+        "default": None
     },
     "profile": {
         "description": "Profile scan performance",

wordfence/cli/malwarescan/malwarescan.py

@@ -7,8 +7,10 @@
 
 from wordfence import scanning
 from wordfence.scanning import filtering
-from wordfence.scanning.matching import MatchEngine, MatchEngineOptions
-from wordfence.util import caching, pcre, serialization
+from wordfence.scanning.matching import MatchEngine, MatchEngineOptions, \
+    MatchEngineCompilerOptions
+from wordfence.util import caching, pcre, serialization, vectorscan
+from wordfence.util.platform import Platform
 from wordfence.intel.signatures import SignatureSet, PrecompiledSignatureSet
 from wordfence.logging import (log, remove_initial_handler,
                                restore_initial_handler)
@@ -58,18 +60,44 @@ def _filter_signatures(
         signature_count = len(signatures.signatures)
         log.debug(f'Filtered signature count: {signature_count}')
 
-    def _get_signatures(self) -> SignatureSet:
+    def _get_pre_compiled_signatures(
+                self,
+                match_engine: MatchEngine
+            ) -> Optional[PrecompiledSignatureSet]:
+
+        def fetch_pre_compiled() -> Optional[PrecompiledSignatureSet]:
+            client = self.context.get_noc1_client()
+            return client.get_precompiled_malware_signatures(
+                    Platform.detect(),
+                    vectorscan.API_VERSION
+                )
+
+        cacheable = caching.Cacheable(
+                'pre-compiled-signatures-{match_engine.module}',
+                fetch_pre_compiled,
+                caching.DURATION_ONE_DAY
+            )
+        return cacheable.get(self.cache)
+
+    def _get_signatures(self, match_engine: MatchEngine) -> SignatureSet:
+        supports_pre_compilation = match_engine.supports_pre_compilation()
+        if supports_pre_compilation:
+            precompiled = self._get_pre_compiled_signatures(match_engine)
+            if precompiled is not None:
+                return precompiled.signature_set, precompiled.data
+
         def fetch_signatures() -> SignatureSet:
             noc1_client = self.context.get_noc1_client()
             return noc1_client.get_malware_signatures()
+
         self.cacheable_signatures = caching.Cacheable(
                 'signatures',
                 fetch_signatures,
                 caching.DURATION_ONE_DAY
             )
         signatures = self.cacheable_signatures.get(self.cache)
         self._filter_signatures(signatures)
-        return signatures
+        return signatures, None
 
     def _get_file_list_separator(self) -> str:
         if isinstance(self.config.file_list_separator, bytes):
@@ -131,13 +159,17 @@ def _get_database_source(
         current_hash = match_engine_options.signature_set.get_hash()
 
         def compile_database():
-            compiler = match_engine.get_compiler(match_engine_options)
+            compiler_options = MatchEngineCompilerOptions(
+                    generic=self.config.pre_compile_generic
+                )
+            compiler = match_engine.get_compiler(compiler_options)
             if compiler is None:
                 return None
             compiled = compiler.compile_serializable(
                     match_engine_options.signature_set
                 )
             return PrecompiledSignatureSet(
+                    match_engine_options.signature_set,
                     current_hash,
                     compiled,
                     match_engine_options.signature_set.license
@@ -146,7 +178,8 @@ def compile_database():
         def is_precompiled_compatible(precompiled):
             return (
                    precompiled is not None and
-                   precompiled.signature_hash == current_hash
+                   precompiled.signature_hash == current_hash and
+                   precompiled.is_supported_version()
                )
 
         if self.config.pattern_database_path is None:
@@ -158,7 +191,7 @@ def filter_precompiled(precompiled):
                         )
                 return precompiled
             cacheable = caching.Cacheable(
-                    f'precompiled-signatures-{match_engine.module}',
+                    f'compiled-signatures-{match_engine.module}',
                     compile_database,
                     filters=[
                         filter_precompiled
@@ -209,20 +242,27 @@ def handle_interrupt(signal_number: int, stack) -> None:
     def invoke(self) -> int:
         self._initialize_interrupt_handling()
 
-        signatures = self._get_signatures()
         match_engine = MatchEngine.for_option(self.config.match_engine)
+        signatures, database_source = self._get_signatures(match_engine)
         match_engine_options = MatchEngineOptions(
                 signature_set=signatures,
                 pcre_options=self._get_pcre_options(),
                 match_all=self.config.match_all
             )
-        match_engine_options.database_source = self._get_database_source(
-                match_engine,
-                match_engine_options,
-                force=self.config.pre_compile
+        pre_compile = (
+                self.config.pre_compile
+                or self.config.pre_compile_generic
+            )
+        match_engine_options.database_source = (
+                database_source if database_source is not None else
+                self._get_database_source(
+                    match_engine,
+                    match_engine_options,
+                    force=pre_compile
+                )
             )
 
-        if self.config.pre_compile:
+        if pre_compile:
             if match_engine_options.database_source is None:
                 log.error(
                         'Signature set pre-compilation is not supported '

wordfence/intel/signatures.py

@@ -1,5 +1,5 @@
 from hashlib import sha256
-from typing import Union
+from typing import Optional
 
 from ..api.licensing import License, LicenseSpecific
 
@@ -81,15 +81,23 @@ def get_hash(self) -> str:
 
 class PrecompiledSignatureSet(LicenseSpecific):
 
+    VERSION = 1
+
     def __init__(
                 self,
-                signature_set: Union[bytes, SignatureSet],
+                signature_set: SignatureSet,
+                signature_hash: Optional[bytes],
                 data: bytes,
                 license: License = None
             ):
         super().__init__(license)
+        self.signature_set = SignatureSet
         self.signature_hash = (
                 signature_set if isinstance(signature_set, bytes)
                 else signature_set.get_hash()
             )
         self.data = data
+        self.version = self.VERSION
+
+    def is_supported_version(self) -> bool:
+        return hasattr(self, 'version') and self.version == self.VERSION

wordfence/scanning/matching/matching.py

@@ -110,6 +110,11 @@ def __init__(self, matcher: Matcher):
         super().__init__()
 
 
+@dataclass
+class MatchEngineCompilerOptions:
+    generic: bool = False
+
+
 @dataclass
 class MatchEngineOptions:
     signature_set: SignatureSet
@@ -164,10 +169,17 @@ def _get_loaded_module(self):
             self._loaded_module = self._load_module()
         return self._loaded_module
 
-    def get_compiler(self, options: MatchEngineOptions) -> Optional[Compiler]:
+    def get_compiler(
+                self,
+                options: MatchEngineCompilerOptions
+            ) -> Optional[Compiler]:
         module = self._get_loaded_module()
         return module.create_compiler(options)
 
+    def supports_pre_compilation(self) -> bool:
+        compiler = self.get_compiler(MatchEngineCompilerOptions())
+        return compiler is not None
+
     def create_matcher(self, options: MatchEngineOptions) -> Matcher:
         module = self._get_loaded_module()
         return module.create_matcher(options)

wordfence/scanning/matching/pcre.py

@@ -8,7 +8,8 @@
         PcreOptions, PCRE_DEFAULT_OPTIONS
 
 from .matching import Matcher, BaseMatcherContext, TimeoutException, \
-        MatchWorkspace, MatchEngineOptions, DEFAULT_TIMEOUT
+        MatchWorkspace, MatchEngineCompilerOptions, MatchEngineOptions, \
+        DEFAULT_TIMEOUT
 
 
 if not pcre.AVAILABLE:
@@ -234,7 +235,7 @@ def create_workspace(self) -> PcreMatchWorkspace:
         return PcreMatchWorkspace()
 
 
-def create_compiler(options: MatchEngineOptions) -> None:
+def create_compiler(options: MatchEngineCompilerOptions) -> None:
     return None
 
 

wordfence/scanning/matching/vectorscan.py

@@ -5,7 +5,7 @@
 from ...util import vectorscan
 
 from .matching import MatchEngineOptions, Matcher, BaseMatcherContext, \
-        MatchWorkspace, Compiler
+        MatchWorkspace, MatchEngineCompilerOptions, Compiler
 
 
 if not vectorscan.AVAILABLE:
@@ -14,7 +14,8 @@
 
 from ...util.vectorscan import VectorscanStreamScanner, VectorscanMatch, \
         VectorscanFlags, VectorscanDatabase, VectorscanScanTerminated, \
-        VectorscanMode, vectorscan_compile, vectorscan_deserialize
+        VectorscanMode, vectorscan_compile, vectorscan_deserialize, \
+        VectorscanPlatformInfo, VectorscanCpuFeatures, VectorscanTuneFamily
 
 
 class VectorscanMatcherContext(BaseMatcherContext):
@@ -54,6 +55,9 @@ def __exit__(self, exc_type, exc_value, traceback) -> None:
 
 class VectorscanCompiler(Compiler):
 
+    def __init__(self, generic: bool = False):
+        self.generic = generic
+
     def compile(self, signature_set: SignatureSet) -> bytes:
         patterns = {
                 signature.identifier: signature.rule
@@ -69,10 +73,15 @@ def compile(self, signature_set: SignatureSet) -> bytes:
                 VectorscanFlags.SINGLEMATCH |
                 VectorscanFlags.ALLOWEMPTY
             )
+        platform_info = VectorscanPlatformInfo(
+                VectorscanCpuFeatures.NONE,
+                VectorscanTuneFamily.GENERIC
+            ) if self.generic else None
         database = vectorscan_compile(
                 patterns,
                 mode=VectorscanMode.STREAM,
-                flags=flags
+                flags=flags,
+                platform_info=platform_info
             )
         log.debug('Successfully compiled vectorscan database')
         return database
@@ -134,8 +143,8 @@ def create_context(self) -> VectorscanMatcherContext:
             )
 
 
-def create_compiler(options: MatchEngineOptions):
-    return VectorscanCompiler()
+def create_compiler(options: MatchEngineCompilerOptions):
+    return VectorscanCompiler(generic=options.generic)
 
 
 def create_matcher(options: MatchEngineOptions) -> VectorscanMatcher: