# SPDX-License-Identifier: CC0-1.0

"""
Haketilo unit tests - enforcing script blocking policy from content script
"""

# This file is part of Haketilo
#
# Copyright (C) 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 json
import urllib.parse
from selenium.webdriver.support.ui import WebDriverWait

from ..script_loader import load_script
from .utils import are_scripts_allowed

# For simplicity, we'll use one nonce in all test cases.
nonce = 'e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855'

allow_policy = {'allow': True}
block_policy = {
    'allow': False,
    'csp': f"prefetch-src 'none'; script-src-attr 'none'; script-src 'none'; script-src-elem 'none'; frame-src http://* https://*;"
}
payload_policy = {
    'mapping': 'somemapping',
    'payload': {'identifier': 'someresource'},
    'csp': f"prefetch-src 'none'; script-src-attr 'none'; script-src 'nonce-{nonce}'; script-src-elem 'nonce-{nonce}';"
}

content_script = load_script('content/policy_enforcing.js') + ''';{
const smuggled_what_to_do = /^[^#]*#?(.*)$/.exec(document.URL)[1];
const what_to_do = smuggled_what_to_do === "" ? {policy: {allow: true}} :
                   JSON.parse(decodeURIComponent(smuggled_what_to_do));

if (what_to_do.csp_off) {
    const orig_DOMParser = window.DOMParser;
    window.DOMParser = function() {
        const parser = new orig_DOMParser();
        this.parseFromString = () => parser.parseFromString('', 'text/html');
    }
}

enforce_blocking(what_to_do.policy);
}'''

def get(driver, page, what_to_do):
    driver.get(page + '#' + urllib.parse.quote(json.dumps(what_to_do)))
    driver.execute_script('window.before_reload = true; location.reload();')
    done = lambda _: not driver.execute_script('return window.before_reload;')
    WebDriverWait(driver, 10).until(done)

@pytest.mark.ext_data({'content_script': content_script})
@pytest.mark.usefixtures('webextension')
# Under Mozilla we use several mechanisms of script blocking. Some serve as
# fallbacks in case others break. CSP one of those mechanisms. Here we run the
# test once with CSP blocking on and once without it. This allows us to verify
# that the CSP-less blocking approaches by themselves also work. We don't do the
# reverse (CSP on and other mechanisms off) because CSP rules added through
# <meta> injection are not reliable enough - they do not always take effect
# immediately and there's nothing we can do to fix it.
@pytest.mark.parametrize('csp_off_setting', [{}, {'csp_off': True}])
def test_policy_enforcing_html(driver, execute_in_page, csp_off_setting):
    """
    A test case of sanitizing <script>s and intrinsic JavaScript in HTML pages.
    """
    def click_all():
        for i in range(1, 3):
            driver.find_element_by_id(f'clickme{i}').click()

    def assert_properly_blocked():
        click_all()

        assert set(driver.execute_script('return window.__run || [];')) == set()
        assert bool(csp_off_setting) == are_scripts_allowed(driver)

        for attr in ('onclick', 'href', 'src', 'data'):
            elem = driver.find_element_by_css_selector(f'[blocked-{attr}]')

            assert 'blocked' in elem.get_attribute(attr)
            assert '__run = [...(' in elem.get_attribute(f'blocked-{attr}')

        but1 = driver.find_element_by_id('clickme1')
        assert but1.get_attribute('blocked-blocked-onclick') == \
            "some useful data"

    # First, see if scripts run when not blocked.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_1.html', {
        'policy': allow_policy,
        **csp_off_setting
    })

    click_all()

    assert set(driver.execute_script('return window.__run || [];')) == \
        {'inline', 'on', 'href', 'src', 'data'}
    assert are_scripts_allowed(driver)

    # Now, verify scripts don't run when blocked.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_1.html', {
        'policy': block_policy,
        **csp_off_setting
    })

    assert_properly_blocked()

    # Now, verify only scripts with nonce can run when payload is injected.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_1.html', {
        'policy': payload_policy,
        **csp_off_setting
    })

    assert_properly_blocked()
    assert are_scripts_allowed(driver, nonce)

# Test function analogous to that for HTML page.
@pytest.mark.ext_data({'content_script': content_script})
@pytest.mark.usefixtures('webextension')
@pytest.mark.parametrize('csp_off_setting', [{}, {'csp_off': True}])
def test_policy_enforcing_xml(driver, execute_in_page, csp_off_setting):
    """
    A test case of sanitizing <script>s and intrinsic JavaScript in XML
    documents.
    """
    def click_all():
        for name in ('idaret', 'nowamak', 'mango', 'annoying'):
            elem = driver.find_element_by_id(f'{name}_circle')
            try:
                elem.click()
            except:
                pass

    def assert_properly_blocked():
        click_all()

        try:
            assert set(driver.execute_script('return window.__run || [];')) == set()
        except:
            from time import sleep
            sleep(100000)
        assert bool(csp_off_setting) == are_scripts_allowed(driver)

    # First, see if scripts run when not blocked.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
        'policy': allow_policy,
        **csp_off_setting
    })

    click_all()

    assert set(driver.execute_script('return window.__run || [];')) == \
        {'grape', 'raspberry', 'idaret', 'melon'}
    assert are_scripts_allowed(driver)

    # Now, verify scripts don't run when blocked.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
        'policy': block_policy,
        **csp_off_setting
    })

    assert_properly_blocked()

    # Now, verify only scripts with nonce can run when payload is injected.
    get(driver, 'https://gotmyowndoma.in/scripts_to_block_2.xml', {
        'policy': payload_policy,
        **csp_off_setting
    })

    assert_properly_blocked()
    assert are_scripts_allowed(driver, nonce)