1#!/usr/bin/env python 2# Copyright 2013 Google Inc. All Rights Reserved. 3# 4# Licensed under the Apache License, Version 2.0 (the "License"); 5# you may not use this file except in compliance with the License. 6# You may obtain a copy of the License at 7# 8# http://www.apache.org/licenses/LICENSE-2.0 9# 10# Unless required by applicable law or agreed to in writing, software 11# distributed under the License is distributed on an "AS IS" BASIS, 12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13# See the License for the specific language governing permissions and 14# limitations under the License. 15 16import datetime 17import httparchive 18import mock 19import script_injector 20import unittest 21 22 23LONG_COMMENT = '<!--' + 'comment,' * 200 + '-->' 24COMMENT_OR_NOT = ('', LONG_COMMENT) 25SCRIPT_TO_INJECT = 'var flag = 0;' 26EXPECTED_SCRIPT = '<script>' + SCRIPT_TO_INJECT + '</script>' 27TEXT_HTML = 'text/html' 28TEXT_CSS = 'text/css' 29APPLICATION = 'application/javascript' 30SEPARATOR = httparchive.ArchivedHttpResponse.CHUNK_EDIT_SEPARATOR 31SEPARATORS_OR_NOT = ('', SEPARATOR, SEPARATOR*3) 32 33TEMPLATE_HEAD = """\ 34{boundary_at_start}\ 35<!doc{boundary_in_doctype}type html>{boundary_after_doctype}\ 36<ht{boundary_in_html}ml>{boundary_after_html}\ 37<he{boundary_in_head}ad>{injection}{boundary_after_head}\ 38</head></html>\ 39""" 40TEMPLATE_HTML = """\ 41{boundary_at_start}\ 42<!doc{boundary_in_doctype}type html>{boundary_after_doctype}\ 43<ht{boundary_in_html}ml>{injection}{boundary_after_html}\ 44</html>\ 45""" 46TEMPLATE_DOCTYPE = """\ 47{boundary_at_start}\ 48<!doc{boundary_in_doctype}type html>{injection}{boundary_after_doctype}\ 49<body></body>\ 50""" 51TEMPLATE_RAW = """\ 52{boundary_at_start}\ 53{injection}<body></body>\ 54""" 55NORMAL_TEMPLATES = (TEMPLATE_HEAD, TEMPLATE_HTML, 56 TEMPLATE_DOCTYPE, TEMPLATE_RAW) 57TEMPLATE_COMMENT = """\ 58{comment_before_doctype}<!doctype html>{comment_after_doctype}\ 59<html>{comment_after_html}<head>{injection}</head></html>\ 60""" 61 62 63def _wrap_inject_script(source, application, script_to_inject): 64 text_chunks = source.split(SEPARATOR) 65 text_chunks, just_injected = script_injector.InjectScript( 66 text_chunks, application, script_to_inject) 67 result = SEPARATOR.join(text_chunks) 68 return result, just_injected 69 70 71class ScriptInjectorTest(unittest.TestCase): 72 73 def _assert_no_injection(self, source, application): 74 new_source, just_injected = _wrap_inject_script( 75 source, application, SCRIPT_TO_INJECT) 76 self.assertEqual(new_source, source) 77 self.assertFalse(just_injected) 78 79 def _assert_successful_injection(self, template): 80 source, just_injected = _wrap_inject_script( 81 template.format(injection=''), TEXT_HTML, SCRIPT_TO_INJECT) 82 self.assertEqual(source, template.format(injection=EXPECTED_SCRIPT)) 83 self.assertTrue(just_injected) 84 85 def test_unsupported_content_type(self): 86 self._assert_no_injection('abc', TEXT_CSS) 87 self._assert_no_injection('abc', APPLICATION) 88 89 def test_empty_content_as_already_injected(self): 90 self._assert_no_injection('', TEXT_HTML) 91 92 def test_non_html_content_with_html_content_type(self): 93 self._assert_no_injection('{"test": 1"}', TEXT_HTML) 94 95 def test_already_injected(self): 96 parameters = {'injection': SCRIPT_TO_INJECT} 97 for template in NORMAL_TEMPLATES: 98 for parameters['boundary_at_start'] in SEPARATORS_OR_NOT: 99 for parameters['boundary_in_doctype'] in SEPARATORS_OR_NOT: 100 for parameters['boundary_after_doctype'] in SEPARATORS_OR_NOT: 101 for parameters['boundary_in_html'] in SEPARATORS_OR_NOT: 102 for parameters['boundary_after_html'] in SEPARATORS_OR_NOT: 103 for parameters['boundary_in_head'] in SEPARATORS_OR_NOT: 104 for parameters['boundary_after_head'] in SEPARATORS_OR_NOT: 105 source = template.format(**parameters) 106 self._assert_no_injection(source, TEXT_HTML) 107 108 def test_normal(self): 109 parameters = {'injection': '{injection}'} 110 for template in NORMAL_TEMPLATES: 111 for parameters['boundary_at_start'] in SEPARATORS_OR_NOT: 112 for parameters['boundary_in_doctype'] in SEPARATORS_OR_NOT: 113 for parameters['boundary_after_doctype'] in SEPARATORS_OR_NOT: 114 for parameters['boundary_in_html'] in SEPARATORS_OR_NOT: 115 for parameters['boundary_after_html'] in SEPARATORS_OR_NOT: 116 for parameters['boundary_in_head'] in SEPARATORS_OR_NOT: 117 for parameters['boundary_after_head'] in SEPARATORS_OR_NOT: 118 template = template.format(**parameters) 119 self._assert_successful_injection(template) 120 121 def test_comments(self): 122 parameters = {'injection': '{injection}'} 123 for parameters['comment_before_doctype'] in COMMENT_OR_NOT: 124 for parameters['comment_after_doctype'] in COMMENT_OR_NOT: 125 for parameters['comment_after_html'] in COMMENT_OR_NOT: 126 template = TEMPLATE_COMMENT.format(**parameters) 127 self._assert_successful_injection(template) 128 129 @mock.patch('script_injector.os.path.exists', return_value=True) 130 @mock.patch('script_injector.open', 131 mock.mock_open(read_data='var time_seed = 123;')) 132 def test_injection_function(self, _): 133 injector = script_injector.GetScriptInjector('to_inject.js') 134 self.assertEqual('var time_seed=123;', 135 injector(datetime.datetime.utcnow())) 136 137 @mock.patch('script_injector.os.path.exists', return_value=True) 138 @mock.patch('script_injector.open', 139 mock.mock_open( 140 read_data='var time_seed = {{WPR_TIME_SEED_TIMESTAMP}};')) 141 def test_time_seed_replacement(self, _): 142 date = datetime.datetime(2016, 11, 17) 143 injector = script_injector.GetScriptInjector('date.js') 144 self.assertEqual('var time_seed=1479340800000;', injector(date)) 145 146 147if __name__ == '__main__': 148 unittest.main() 149