summaryrefslogtreecommitdiffstats
path: root/testing/firefox-ui/tests/functional/safebrowsing/test_initial_download.py
blob: 4013c51d5bffb751655efed13f79c9a654b1e0ce (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

import os
from functools import reduce

from marionette_driver import Wait
from marionette_harness import MarionetteTestCase


class TestSafeBrowsingInitialDownload(MarionetteTestCase):
    v2_file_extensions = [
        "vlpset",
        "sbstore",
    ]

    v4_file_extensions = [
        "vlpset",
        "metadata",
    ]

    prefs_download_lists = [
        "urlclassifier.blockedTable",
        "urlclassifier.downloadAllowTable",
        "urlclassifier.downloadBlockTable",
        "urlclassifier.malwareTable",
        "urlclassifier.phishTable",
        "urlclassifier.trackingTable",
        "urlclassifier.trackingWhitelistTable",
    ]

    prefs_provider_update_time = {
        # Force an immediate download of the safebrowsing files
        "browser.safebrowsing.provider.mozilla.nextupdatetime": 1,
    }

    prefs_safebrowsing = {
        "browser.safebrowsing.debug": True,
        "browser.safebrowsing.update.enabled": True,
    }

    def get_safebrowsing_files(self, is_v4):
        files = []

        if is_v4:
            my_file_extensions = self.v4_file_extensions
        else:  # v2
            my_file_extensions = self.v2_file_extensions

        for pref_name in self.prefs_download_lists:
            base_names = self.marionette.get_pref(pref_name).split(",")

            # moztest- lists are not saved to disk
            # pylint --py3k: W1639
            base_names = list(
                filter(lambda x: not x.startswith("moztest-"), base_names)
            )

            for ext in my_file_extensions:
                files.extend(
                    [
                        "{name}.{ext}".format(name=f, ext=ext)
                        for f in base_names
                        if f and f.endswith("-proto") == is_v4
                    ]
                )

        return set(sorted(files))

    def setUp(self):
        super(TestSafeBrowsingInitialDownload, self).setUp()

        self.safebrowsing_v2_files = self.get_safebrowsing_files(False)
        if any(
            f.startswith("goog-") or f.startswith("googpub-")
            for f in self.safebrowsing_v2_files
        ):
            self.prefs_provider_update_time.update(
                {
                    "browser.safebrowsing.provider.google.nextupdatetime": 1,
                }
            )

        self.safebrowsing_v4_files = self.get_safebrowsing_files(True)
        if any(
            f.startswith("goog-") or f.startswith("googpub-")
            for f in self.safebrowsing_v4_files
        ):
            self.prefs_provider_update_time.update(
                {
                    "browser.safebrowsing.provider.google4.nextupdatetime": 1,
                }
            )

        # Force the preferences for the new profile
        enforce_prefs = self.prefs_safebrowsing
        enforce_prefs.update(self.prefs_provider_update_time)
        self.marionette.enforce_gecko_prefs(enforce_prefs)

        self.safebrowsing_path = os.path.join(
            self.marionette.instance.profile.profile, "safebrowsing"
        )

    def tearDown(self):
        try:
            # Restart with a fresh profile
            self.marionette.restart(in_app=False, clean=True)
        finally:
            super(TestSafeBrowsingInitialDownload, self).tearDown()

    def test_safe_browsing_initial_download(self):
        def check_downloaded(_):
            return reduce(
                lambda state, pref: state and int(self.marionette.get_pref(pref)) != 1,
                list(self.prefs_provider_update_time),
                True,
            )

        try:
            Wait(self.marionette, timeout=170).until(
                check_downloaded,
                message="Not all safebrowsing files have been downloaded",
            )
        finally:
            files_on_disk_toplevel = os.listdir(self.safebrowsing_path)
            for f in self.safebrowsing_v2_files:
                self.assertIn(f, files_on_disk_toplevel)

            if len(self.safebrowsing_v4_files) > 0:
                files_on_disk_google4 = os.listdir(
                    os.path.join(self.safebrowsing_path, "google4")
                )
                for f in self.safebrowsing_v4_files:
                    self.assertIn(f, files_on_disk_google4)