Create a run_yara task (#12005)

.gitignore

@@ -46,3 +46,6 @@ storage/guarded-addons/*
 storage/shared_storage/*
 supervisord.pid
 tmp/*
+
+# private projects
+addons-yara

requirements/prod.txt

@@ -391,3 +391,17 @@ atpublic==1.0 \
     --hash=sha256:7dca670499e9a9d3aae5a8914bc799475fe24be3bcd29c8129642dda665f7a44
 flufl.lock==3.2 \
     --hash=sha256:a8d66accc9ab41f09961cd8f8db39f9c28e97e2769659a3567c63930a869ff5b
+yara-python==3.10.0 \
+    --hash=sha256:024fb818ece339a378c562b3b49c28bf861b004ab5c57e488556d49bcd8ab76d \
+    --hash=sha256:04a2a9e31668aa20e8ad4e3cd8273b9216391f0e98704ce6637758f644837f78 \
+    --hash=sha256:0a81aef159adf9a233bf52168cb7ed63edd3213162c1e88aa234a252b46efcf7 \
+    --hash=sha256:0bf33b91e0dacaff8590631f6459fdb084271de4cb1d0ecbeff2401f9ae20d5c \
+    --hash=sha256:16bbd479078f3cce33b4966c112913c0c9a7d22d0a63df23e06be5a84e51d5f4 \
+    --hash=sha256:2da1d94850cbea1dd9db1cc7d54bb36a69cd6a33bbc0caf003497b6a323e3e10 \
+    --hash=sha256:6266b2798a6d8df34182658e00556de9a73fcd5fa736a936a4184a0d7dffa41e \
+    --hash=sha256:71abafd7ae3fe30d2c5beb3ccdd6cb3a7f9a6f4edbc2d8f120ad40d4ff35383a \
+    --hash=sha256:759e11705345d0bc429bedb59ce63ccdab1959f5ec44f6fbc1a24a59f224fd30 \
+    --hash=sha256:7980282496d73d5854ead8d0cb134b1437bf1c6bf920bbe8dbafcaca0dc8a482 \
+    --hash=sha256:c5385f28efd42ec615870f30572576b6e3cc12d08025178c9040979244e42eb3 \
+    --hash=sha256:cd4b6da645a210dab5beff948effbab2861d5620b2fcd3724f1182870154ca4d \
+    --hash=sha256:fb74b068109ec359197abcabfe5d777e36aa7c42ed9bdec3c93de1c40c215b03

src/olympia/lib/settings_base.py

@@ -1628,7 +1628,6 @@ STATICFILES_DIRS = (
 # Make sure to check overwrites in conftest.py if new settings are added
 # or changed.
 STORAGE_ROOT = env('NETAPP_STORAGE_ROOT', default=path('storage'))
-
 ADDONS_PATH = os.path.join(STORAGE_ROOT, 'files')
 GUARDED_ADDONS_PATH = os.path.join(STORAGE_ROOT, 'guarded-addons')
 GIT_FILE_STORAGE_PATH = os.path.join(STORAGE_ROOT, 'git-storage')
@@ -1638,6 +1637,9 @@ SHARED_STORAGE = os.path.join(STORAGE_ROOT, 'shared_storage')
 MEDIA_ROOT = os.path.join(SHARED_STORAGE, 'uploads')
 TMP_PATH = os.path.join(SHARED_STORAGE, 'tmp')
 
+YARA_ROOT = env('YARA_ROOT', default=path('addons-yara'))
+YARA_RULES_FILEPATH = os.path.join(YARA_ROOT, 'rules', 'all_rules.yar')
+
 # These are key files that must be present on disk to encrypt/decrypt certain
 # database fields.
 # {'api_key:secret': os.path.join(ROOT, 'path', 'to', 'file.key'),}

src/olympia/yara/tasks.py

@@ -0,0 +1,59 @@
+import yara
+
+from django.conf import settings
+
+import olympia.core.logger
+
+from olympia.amo.celery import task
+from olympia.files.models import FileUpload
+from olympia.files.utils import SafeZip
+
+from .models import YaraResult
+
+log = olympia.core.logger.getLogger('z.yara.task')
+
+
+@task
+def run_yara(upload_pk):
+    """
+    Apply a set of Yara rules on a FileUpload and store the results.
+
+    This task is intended to be run as part of the submission process only.
+    When a version is created from a FileUpload, the files are removed. In
+    addition, we usually delete old FileUpload entries after 180 days.
+    """
+    log.info('Starting yara task for FileUpload %s.', upload_pk)
+    upload = FileUpload.objects.get(pk=upload_pk)
+
+    if not upload.path.endswith('.xpi'):
+        log.info('Not running yara for FileUpload %s, it is not a xpi file.',
+                 upload_pk)
+        return
+
+    try:
+        rules = yara.compile(filepath=settings.YARA_RULES_FILEPATH)
+
+        result = YaraResult()
+        result.upload = upload
+
+        zip_file = SafeZip(source=upload.path)
+        for zip_info in zip_file.info_list:
+            if not zip_info.is_dir():
+                file_content = zip_file.read(zip_info).decode(errors='ignore')
+                for match in rules.match(data=file_content):
+                    # Add the filename to the meta dict.
+                    meta = {**match.meta, 'filename': zip_info.filename}
+                    result.add_match(
+                        rule=match.rule,
+                        tags=match.tags,
+                        meta=meta
+                    )
+
+        zip_file.close()
+        result.save()
+
+        log.info('Ending yara task for FileUpload %s.', upload_pk)
+    except Exception:
+        # We log the exception but we do not raise to avoid perturbing the
+        # submission flow.
+        log.exception('Error in yara task for FileUpload %s.', upload_pk)

src/olympia/yara/tests/test_task.py

@@ -0,0 +1,87 @@
+import yara
+
+from unittest import mock
+
+from django.test.utils import override_settings
+
+from olympia.amo.tests import TestCase
+from olympia.files.tests.test_models import UploadTest
+from olympia.yara.models import YaraResult
+from olympia.yara.tasks import run_yara
+
+
+class TestRunYara(UploadTest, TestCase):
+    @mock.patch('yara.compile')
+    def test_skip_non_xpi_files_with_mocks(self, yara_compile_mock):
+        upload = self.get_upload('search.xml')
+
+        run_yara(upload.pk)
+
+        assert not yara_compile_mock.called
+
+    def test_run_with_mocks(self):
+        upload = self.get_upload('webextension.xpi')
+        assert len(YaraResult.objects.all()) == 0
+
+        # This compiled rule will match for all files in the xpi.
+        rules = yara.compile(source='rule always_true { condition: true }')
+        with mock.patch('yara.compile') as yara_compile_mock:
+            yara_compile_mock.return_value = rules
+            run_yara(upload.pk)
+
+        assert yara_compile_mock.called
+        results = YaraResult.objects.all()
+        assert len(results) == 1
+        result = results[0]
+        assert result.upload == upload
+        assert len(result.matches) == 2
+        assert result.matches[0] == {
+            'rule': 'always_true',
+            'tags': [],
+            'meta': {
+                'filename': 'index.js'
+            },
+        }
+        assert result.matches[1] == {
+            'rule': 'always_true',
+            'tags': [],
+            'meta': {
+                'filename': 'manifest.json'
+            },
+        }
+
+    def test_run_no_matches_with_mocks(self):
+        upload = self.get_upload('webextension.xpi')
+        assert len(YaraResult.objects.all()) == 0
+
+        # This compiled rule will never match.
+        rules = yara.compile(source='rule always_false { condition: false }')
+        with mock.patch('yara.compile') as yara_compile_mock:
+            yara_compile_mock.return_value = rules
+            run_yara(upload.pk)
+
+        result = YaraResult.objects.all()[0]
+        assert result.matches == []
+
+    def test_run_ignores_directories(self):
+        upload = self.get_upload('webextension_signed_already.xpi')
+        # This compiled rule will match for all files in the xpi.
+        rules = yara.compile(source='rule always_true { condition: true }')
+
+        with mock.patch('yara.compile') as yara_compile_mock:
+            yara_compile_mock.return_value = rules
+            run_yara(upload.pk)
+
+        result = YaraResult.objects.all()[0]
+        assert result.upload == upload
+        # The `webextension_signed_already.xpi` fixture file has 1 directory
+        # and 3 files.
+        assert len(result.matches) == 3
+
+    @override_settings(YARA_RULES_FILEPATH='unknown/path/to/rules.yar')
+    def test_run_does_not_raise(self):
+        upload = self.get_upload('webextension.xpi')
+
+        # This call should not raise even though there will be an error because
+        # YARA_RULES_FILEPATH is configured with a wrong path.
+        run_yara(upload.pk)