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 16 import datetime 17 import httparchive 18 import mock 19 import script_injector 20 import unittest 21 22 23 LONG_COMMENT = '<!--' + 'comment,' * 200 + '-->' 24 COMMENT_OR_NOT = ('', LONG_COMMENT) 25 SCRIPT_TO_INJECT = 'var flag = 0;' 26 EXPECTED_SCRIPT = '<script>' + SCRIPT_TO_INJECT + '</script>' 27 TEXT_HTML = 'text/html' 28 TEXT_CSS = 'text/css' 29 APPLICATION = 'application/javascript' 30 SEPARATOR = httparchive.ArchivedHttpResponse.CHUNK_EDIT_SEPARATOR 31 SEPARATORS_OR_NOT = ('', SEPARATOR, SEPARATOR*3) 32 33 TEMPLATE_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 """ 40 TEMPLATE_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 """ 46 TEMPLATE_DOCTYPE = """\ 47 {boundary_at_start}\ 48 <!doc{boundary_in_doctype}type html>{injection}{boundary_after_doctype}\ 49 <body></body>\ 50 """ 51 TEMPLATE_RAW = """\ 52 {boundary_at_start}\ 53 {injection}<body></body>\ 54 """ 55 NORMAL_TEMPLATES = (TEMPLATE_HEAD, TEMPLATE_HTML, 56 TEMPLATE_DOCTYPE, TEMPLATE_RAW) 57 TEMPLATE_COMMENT = """\ 58 {comment_before_doctype}<!doctype html>{comment_after_doctype}\ 59 <html>{comment_after_html}<head>{injection}</head></html>\ 60 """ 61 62 63 def _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 71 class 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 147 if __name__ == '__main__': 148 unittest.main() 149
