path: root/test/unit/test_patterns_query_manager.py

# SPDX-License-Identifier: CC0-1.0

"""
Haketilo unit tests - building pattern tree and putting it in a content script
"""

# This file is part of Haketilo
#
# Copyright (C) 2021,2022 Wojtek Kosior <koszko@koszko.org>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the CC0 1.0 Universal License as published by
# the Creative Commons Corporation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# CC0 1.0 Universal License for more details.

import pytest
import re
import json
from selenium.webdriver.support.ui import WebDriverWait

from ..script_loader import load_script

def simple_sample_mapping(patterns, fruit):
    if type(patterns) is not list:
        patterns = [patterns]
    payloads = dict([(p, {'identifier': f'{fruit}-{p}'}) for p in patterns])
    return {
        'source_copyright': [],
        'type': 'mapping',
        'identifier': f'inject-{fruit}',
        'payloads': payloads
    }

content_script_tree_re = re.compile(r'this.haketilo_pattern_tree = (.*);')
def extract_tree_data(content_script_text):
    return json.loads(content_script_tree_re.search(content_script_text)[1])

content_script_mapping_re = re.compile(r'this.haketilo_mappings = (.*);')
def extract_mappings_data(content_script_text):
    return json.loads(content_script_mapping_re.search(content_script_text)[1])

# Fields that are not relevant for testing are omitted from these mapping
# definitions.
sample_mappings = [simple_sample_mapping(pats, fruit) for pats, fruit in [
    (['https://gotmyowndoma.in/index.html',
      'http://gotmyowndoma.in/index.html'], 'banana'),
    (['https://***.gotmyowndoma.in/index.html',
      'https://**.gotmyowndoma.in/index.html',
      'https://*.gotmyowndoma.in/index.html',
      'https://gotmyowndoma.in/index.html'], 'orange'),
    ('https://gotmyowndoma.in/index.html/***', 'grape'),
    ('http://gotmyowndoma.in/index.html/***', 'melon'),
    ('https://gotmyowndoma.in/index.html', 'peach'),
    ('https://gotmyowndoma.in/*', 'pear'),
    ('https://gotmyowndoma.in/**', 'raspberry'),
    ('https://gotmyowndoma.in/***', 'strawberry'),
    ('https://***.gotmyowndoma.in/index.html', 'apple'),
    ('https://***.gotmyowndoma.in/*', 'avocado'),
    ('https://***.gotmyowndoma.in/**', 'papaya'),
    ('https://***.gotmyowndoma.in/***', 'kiwi')
]]

sample_blocking = [f'http{s}://{dw}gotmyown%sdoma.in{i}{pw}'
                   for dw in ('', '***.', '**.', '*.')
                   for i in ('/index.html', '')
                   for pw in ('', '/', '/*')
                   for s in ('', 's')]
sample_blocking = [{'pattern': pattern % (i if i > 1 else ''),
                    'allow': bool(i & 1)}
                   for i, pattern in enumerate(sample_blocking)]

# Even though patterns_query_manager.js is normally meant to run from background
# page, some tests can be as well performed running it from a normal page.
@pytest.mark.get_page('https://gotmyowndoma.in')
def test_pqm_tree_building(driver, execute_in_page):
    """
    patterns_query_manager.js tracks Haketilo's internal database and builds a
    constantly-updated pattern tree based on its contents. Mock the database and
    verify tree building works properly.
    """
    execute_in_page(load_script('background/patterns_query_manager.js'))
    # Mock IndexedDB and build patterns tree.
    execute_in_page(
        '''
        const [initial_mappings, initial_blocking] = arguments.slice(0, 2);
        let mappingchange, blockingchange;

        haketilodb.track.mappings = function (cb) {
            mappingchange = cb;

            return [{}, initial_mappings];
        }
        haketilodb.track.blocking = function (cb) {
            blockingchange = cb;

            return [{}, initial_blocking];
        }

        let last_script;
        let unregister_called = 0;
        async function register_mock(injection)
        {
            await new Promise(resolve => setTimeout(resolve, 1));
            last_script = injection.js[0].code;
            return {unregister: () => unregister_called++};
        }
        browser = {contentScripts: {register: register_mock}};

        returnval(start());
        ''',
        sample_mappings[0:2], sample_blocking[0:2])

    found, tree, content_script, deregistrations = execute_in_page(
        '''
        returnval([pqt.search(tree, arguments[0]).next().value,
                   tree, last_script, unregister_called]);
        ''',
        'https://gotmyowndoma.in/index.html')
    best_pattern = 'https://gotmyowndoma.in/index.html'
    assert found == \
        dict([('~allow', 1),
              *[(f'inject-{fruit}', {'identifier': f'{fruit}-{best_pattern}'})
                for fruit in ('banana', 'orange')]])
    assert tree == extract_tree_data(content_script)
    assert deregistrations == 0

    def condition_all_added(driver):
        last_script = execute_in_page('returnval(last_script);')
        nums = [i for i in range(len(sample_blocking)) if i > 1]
        return (all([('gotmyown%sdoma' % i) in last_script for i in nums]) and
                all([m['identifier'] in last_script for m in sample_mappings]))

    execute_in_page(
        '''
        for (const mapping of arguments[0])
            mappingchange({key: mapping.identifier, new_val: mapping});
        for (const blocking of arguments[1])
            blockingchange({key: blocking.pattern, new_val: blocking});
        ''',
        sample_mappings[2:], sample_blocking[2:])
    WebDriverWait(driver, 10).until(condition_all_added)

    odd_mappings = \
        [m['identifier'] for i, m in enumerate(sample_mappings) if i & 1]
    odd_blocking = \
        [b['pattern'] for i, b in enumerate(sample_blocking) if i & 1]
    even_mappings = \
        [m['identifier'] for i, m in enumerate(sample_mappings) if 1 - i & 1]
    even_blocking = \
        [b['pattern'] for i, b in enumerate(sample_blocking) if 1 - i & 1]

    def condition_odd_removed(driver):
        last_script = execute_in_page('returnval(last_script);')
        nums = [i for i in range(len(sample_blocking)) if i > 1 and 1 - i & 1]
        return (all([id not in last_script for id in odd_mappings]) and
                all([id in last_script for id in even_mappings]) and
                all([p not in last_script for p in odd_blocking[1:]]) and
                all([('gotmyown%sdoma' % i) in last_script for i in nums]))

    def condition_all_removed(driver):
        content_script = execute_in_page('returnval(last_script);')
        return extract_tree_data(content_script) == {}

    execute_in_page(
        '''
        arguments[0].forEach(identifier => mappingchange({key: identifier}));
        arguments[1].forEach(pattern => blockingchange({key: pattern}));
        ''',
        odd_mappings, odd_blocking)

    WebDriverWait(driver, 10).until(condition_odd_removed)

    execute_in_page(
        '''
        arguments[0].forEach(identifier => mappingchange({key: identifier}));
        arguments[1].forEach(pattern => blockingchange({key: pattern}));
        ''',
        even_mappings, even_blocking)

    WebDriverWait(driver, 10).until(condition_all_removed)

content_js = '''
let already_run = false;
this.haketilo_content_script_main = function() {
    if (already_run)
        return;
    already_run = true;
    document.documentElement.innerHTML = "<body><div id='tree-json'>";
    document.getElementById("tree-json").innerText =
        JSON.stringify(this.haketilo_pattern_tree);
}
if (this.haketilo_pattern_tree !== undefined)
    this.haketilo_content_script_main();
'''

def background_js():
    pqm_js = load_script('background/patterns_query_manager.js',
                         "#IMPORT background/broadcast_broker.js")
    return pqm_js + '; broadcast_broker.start(); start();'

@pytest.mark.ext_data({
    'content_script':    content_js,
    'background_script': background_js
})
@pytest.mark.usefixtures('webextension')
def test_pqm_script_injection(driver, execute_in_page):
    # Let's open a normal page in a second window. Window 0 will be used to make
    # changed to IndexedDB and window 1 to test the working of content scripts.
    driver.execute_script('window.open("about:blank", "_blank");')
    windows = [*driver.window_handles]
    assert len(windows) == 2

    def run_content_script():
        driver.switch_to.window(windows[1])
        driver.get('https://gotmyowndoma.in/index.html')
        windows[1] = driver.window_handles[1]
        return driver.execute_script(
            '''
            return (document.getElementById("tree-json") || {}).innerText;
            ''')

    for attempt in range(10):
        json_txt = run_content_script()
        if json.loads(json_txt) == {}:
            break;
        assert attempt != 9

    driver.switch_to.window(windows[0])
    execute_in_page(load_script('common/indexeddb.js'))

    sample_data = {
        'mappings': dict([(sm['identifier'], {'1.0': sm})
                          for sm in sample_mappings]),
        'resources': {},
        'files': {}
    }
    execute_in_page('returnval(save_items(arguments[0]));', sample_data)

    for attempt in range(10):
        tree_json = run_content_script()
        json.loads(tree_json)
        if all([m['identifier'] in tree_json for m in sample_mappings]):
            break
        assert attempt != 9

    driver.switch_to.window(windows[0])
    execute_in_page(
        '''{
        const identifiers = arguments[0];
        async function remove_items()
        {
            const ctx = await start_items_transaction(["mappings"], {});
            for (const id of identifiers)
                await remove_mapping(id, ctx);
            await finalize_transaction(ctx);
        }
        returnval(remove_items());
        }''',
        [sm['identifier'] for sm in sample_mappings])

    for attempt in range(10):
        if json.loads(run_content_script()) == {}:
            break
        assert attempt != 9