get_file_by_url, is_valid_domain_format,
move_hosts_file_into_place, normalize_rule,
path_join_robust, print_failure, print_success,
- supports_color, query_yes_no, recursive_glob,
- remove_old_hosts_file, strip_rule,
+ prompt_for_exclusions, prompt_for_move,
+ prompt_for_flush_dns_cache, prompt_for_update,
+ query_yes_no, recursive_glob,
+ remove_old_hosts_file, supports_color, strip_rule,
update_readme_data, write_data,
write_opening_header)
import updateHostsFile
import mock
-# Base Test Classes
+# Test Helper Objects
class Base(unittest.TestCase):
@staticmethod
def tearDown(self):
shutil.rmtree(self.test_dir)
-# End Base Test Classes
+
+
+def builtins():
+ if PY3:
+ return "builtins"
+ else:
+ return "__builtin__"
+# End Test Helper Objects
# Project Settings
# End Project Settings
+# Prompt the User
+class TestPromptForUpdate(BaseStdout, BaseMockDir):
+
+ def setUp(self):
+ BaseStdout.setUp(self)
+ BaseMockDir.setUp(self)
+
+ def test_no_freshen_no_new_file(self):
+ hosts_file = os.path.join(self.test_dir, "hosts")
+ hosts_data = "This data should not be overwritten"
+
+ with self.mock_property("updateHostsFile.BASEDIR_PATH"):
+ updateHostsFile.BASEDIR_PATH = self.test_dir
+
+ with open(hosts_file, "w") as f:
+ f.write(hosts_data)
+
+ for update_auto in (False, True):
+ dir_count = self.dir_count
+ prompt_for_update(freshen=False, update_auto=update_auto)
+
+ output = sys.stdout.getvalue()
+ self.assertEqual(output, "")
+
+ sys.stdout = StringIO()
+
+ self.assertEqual(self.dir_count, dir_count)
+
+ with open(hosts_file, "r") as f:
+ contents = f.read()
+ self.assertEqual(contents, hosts_data)
+
+ def test_no_freshen_new_file(self):
+ hosts_file = os.path.join(self.test_dir, "hosts")
+
+ with self.mock_property("updateHostsFile.BASEDIR_PATH"):
+ updateHostsFile.BASEDIR_PATH = self.test_dir
+
+ dir_count = self.dir_count
+ prompt_for_update(freshen=False, update_auto=False)
+
+ output = sys.stdout.getvalue()
+ self.assertEqual(output, "")
+
+ sys.stdout = StringIO()
+
+ self.assertEqual(self.dir_count, dir_count + 1)
+
+ with open(hosts_file, "r") as f:
+ contents = f.read()
+ self.assertEqual(contents, "")
+
+ @mock.patch(builtins() + ".open")
+ def test_no_freshen_fail_new_file(self, mock_open):
+ for exc in (IOError, OSError):
+ mock_open.side_effect = exc("failed open")
+
+ with self.mock_property("updateHostsFile.BASEDIR_PATH"):
+ updateHostsFile.BASEDIR_PATH = self.test_dir
+ prompt_for_update(freshen=False, update_auto=False)
+
+ output = sys.stdout.getvalue()
+ expected = ("ERROR: No 'hosts' file in the folder. "
+ "Try creating one manually.")
+ self.assertIn(expected, output)
+
+ sys.stdout = StringIO()
+
+ @mock.patch("updateHostsFile.update_all_sources", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def test_freshen_no_update(self, _, mock_update):
+ hosts_file = os.path.join(self.test_dir, "hosts")
+ hosts_data = "This data should not be overwritten"
+
+ with self.mock_property("updateHostsFile.BASEDIR_PATH"):
+ updateHostsFile.BASEDIR_PATH = self.test_dir
+
+ with open(hosts_file, "w") as f:
+ f.write(hosts_data)
+
+ dir_count = self.dir_count
+
+ prompt_for_update(freshen=True, update_auto=False)
+
+ mock_update.assert_not_called()
+ mock_update.reset_mock()
+
+ output = sys.stdout.getvalue()
+ expected = ("OK, we'll stick with "
+ "what we've got locally.")
+ self.assertIn(expected, output)
+
+ sys.stdout = StringIO()
+
+ self.assertEqual(self.dir_count, dir_count)
+
+ with open(hosts_file, "r") as f:
+ contents = f.read()
+ self.assertEqual(contents, hosts_data)
+
+ @mock.patch("updateHostsFile.update_all_sources", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=True)
+ def test_freshen_update(self, _, mock_update):
+ hosts_file = os.path.join(self.test_dir, "hosts")
+ hosts_data = "This data should not be overwritten"
+
+ with self.mock_property("updateHostsFile.BASEDIR_PATH"):
+ updateHostsFile.BASEDIR_PATH = self.test_dir
+
+ with open(hosts_file, "w") as f:
+ f.write(hosts_data)
+
+ dir_count = self.dir_count
+
+ for update_auto in (False, True):
+ prompt_for_update(freshen=True, update_auto=update_auto)
+
+ mock_update.assert_called_once()
+ mock_update.reset_mock()
+
+ output = sys.stdout.getvalue()
+ self.assertEqual(output, "")
+
+ sys.stdout = StringIO()
+
+ self.assertEqual(self.dir_count, dir_count)
+
+ with open(hosts_file, "r") as f:
+ contents = f.read()
+ self.assertEqual(contents, hosts_data)
+
+ def tearDown(self):
+ BaseStdout.tearDown(self)
+ BaseStdout.tearDown(self)
+
+
+class TestPromptForExclusions(BaseStdout):
+
+ @mock.patch("updateHostsFile.display_exclusion_options", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testSkipPrompt(self, mock_query, mock_display):
+ prompt_for_exclusions(skip_prompt=True)
+
+ output = sys.stdout.getvalue()
+ self.assertEqual(output, "")
+
+ mock_query.assert_not_called()
+ mock_display.assert_not_called()
+
+ @mock.patch("updateHostsFile.display_exclusion_options", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testNoSkipPromptNoDisplay(self, mock_query, mock_display):
+ prompt_for_exclusions(skip_prompt=False)
+
+ output = sys.stdout.getvalue()
+ expected = ("OK, we'll only exclude "
+ "domains in the whitelist.")
+ self.assertIn(expected, output)
+
+ mock_query.assert_called_once()
+ mock_display.assert_not_called()
+
+ @mock.patch("updateHostsFile.display_exclusion_options", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=True)
+ def testNoSkipPromptDisplay(self, mock_query, mock_display):
+ prompt_for_exclusions(skip_prompt=False)
+
+ output = sys.stdout.getvalue()
+ self.assertEqual(output, "")
+
+ mock_query.assert_called_once()
+ mock_display.assert_called_once()
+
+
+class TestPromptForFlushDnsCache(Base):
+
+ @mock.patch("updateHostsFile.flush_dns_cache", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testFlushCache(self, mock_query, mock_flush):
+ for prompt_flush in (False, True):
+ prompt_for_flush_dns_cache(flush_cache=True,
+ prompt_flush=prompt_flush)
+
+ mock_query.assert_not_called()
+ mock_flush.assert_called_once()
+
+ mock_query.reset_mock()
+ mock_flush.reset_mock()
+
+ @mock.patch("updateHostsFile.flush_dns_cache", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testNoFlushCacheNoPrompt(self, mock_query, mock_flush):
+ prompt_for_flush_dns_cache(flush_cache=False,
+ prompt_flush=False)
+
+ mock_query.assert_not_called()
+ mock_flush.assert_not_called()
+
+ @mock.patch("updateHostsFile.flush_dns_cache", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testNoFlushCachePromptNoFlush(self, mock_query, mock_flush):
+ prompt_for_flush_dns_cache(flush_cache=False,
+ prompt_flush=True)
+
+ mock_query.assert_called_once()
+ mock_flush.assert_not_called()
+
+ @mock.patch("updateHostsFile.flush_dns_cache", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=True)
+ def testNoFlushCachePromptFlush(self, mock_query, mock_flush):
+ prompt_for_flush_dns_cache(flush_cache=False,
+ prompt_flush=True)
+
+ mock_query.assert_called_once()
+ mock_flush.assert_called_once()
+
+
+class TestPromptForMove(Base):
+
+ def setUp(self):
+ Base.setUp(self)
+ self.final_file = "final.txt"
+
+ def prompt_for_move(self, **move_params):
+ return prompt_for_move(self.final_file, **move_params)
+
+ @mock.patch("updateHostsFile.move_hosts_file_into_place", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testSkipStaticHosts(self, mock_query, mock_move):
+ for replace in (False, True):
+ for auto in (False, True):
+ move_file = self.prompt_for_move(replace=replace, auto=auto,
+ skipstatichosts=True)
+ self.assertFalse(move_file)
+
+ mock_query.assert_not_called()
+ mock_move.assert_not_called()
+
+ mock_query.reset_mock()
+ mock_move.reset_mock()
+
+ @mock.patch("updateHostsFile.move_hosts_file_into_place", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testReplaceNoSkipStaticHosts(self, mock_query, mock_move):
+ for auto in (False, True):
+ move_file = self.prompt_for_move(replace=True, auto=auto,
+ skipstatichosts=False)
+ self.assertTrue(move_file)
+
+ mock_query.assert_not_called()
+ mock_move.assert_called_once()
+
+ mock_query.reset_mock()
+ mock_move.reset_mock()
+
+ @mock.patch("updateHostsFile.move_hosts_file_into_place", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testAutoNoSkipStaticHosts(self, mock_query, mock_move):
+ for replace in (False, True):
+ move_file = self.prompt_for_move(replace=replace, auto=True,
+ skipstatichosts=True)
+ self.assertFalse(move_file)
+
+ mock_query.assert_not_called()
+ mock_move.assert_not_called()
+
+ mock_query.reset_mock()
+ mock_move.reset_mock()
+
+ @mock.patch("updateHostsFile.move_hosts_file_into_place", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=False)
+ def testPromptNoMove(self, mock_query, mock_move):
+ move_file = self.prompt_for_move(replace=False, auto=False,
+ skipstatichosts=False)
+ self.assertFalse(move_file)
+
+ mock_query.assert_called_once()
+ mock_move.assert_not_called()
+
+ mock_query.reset_mock()
+ mock_move.reset_mock()
+
+ @mock.patch("updateHostsFile.move_hosts_file_into_place", return_value=0)
+ @mock.patch("updateHostsFile.query_yes_no", return_value=True)
+ def testPromptMove(self, mock_query, mock_move):
+ move_file = self.prompt_for_move(replace=False, auto=False,
+ skipstatichosts=False)
+ self.assertTrue(move_file)
+
+ mock_query.assert_called_once()
+ mock_move.assert_called_once()
+
+ mock_query.reset_mock()
+ mock_move.reset_mock()
+# End Prompt the User
+
+
# Exclusion Logic
class TestGatherCustomExclusions(BaseStdout):
settings["extensions"] = sorted(list(
set(options["extensions"]).intersection(settings["extensions"])))
- prompt_for_update()
- prompt_for_exclusions()
+ auto = settings["auto"]
+
+ prompt_for_update(freshen=settings["freshen"], update_auto=auto)
+ prompt_for_exclusions(skip_prompt=auto)
merge_file = create_initial_file()
remove_old_hosts_file(settings["backup"])
final_file = remove_dups_and_excl(merge_file)
number_of_rules = settings["numberofrules"]
+ skip_static_hosts = settings["skipstatichosts"]
write_opening_header(final_file, extensions=extensions,
numberofrules=number_of_rules,
outputsubfolder=output_subfolder,
- skipstatichosts=settings["skipstatichosts"])
+ skipstatichosts=skip_static_hosts)
final_file.close()
if settings["ziphosts"]:
"{:,}".format(number_of_rules) +
" unique entries.")
- prompt_for_move(final_file)
+ move_file = prompt_for_move(final_file, auto=auto,
+ replace=settings["replace"],
+ skipstatichosts=skip_static_hosts)
+
+ # We only flush the DNS cache if we have
+ # moved a new hosts file into place.
+ if move_file:
+ prompt_for_flush_dns_cache(flush_cache=settings["flushdnscache"],
+ prompt_flush=not auto)
# Prompt the User
-def prompt_for_update():
+def prompt_for_update(freshen, update_auto):
"""
Prompt the user to update all hosts files.
+
+ If requested, the function will update all data sources after it
+ checks that a hosts file does indeed exist.
+
+ Parameters
+ ----------
+ freshen : bool
+ Whether data sources should be updated. This function will return
+ if it is requested that data sources not be updated.
+ update_auto : bool
+ Whether or not to automatically update all data sources.
"""
- # Create hosts file if it doesn't exist.
- if not os.path.isfile(path_join_robust(BASEDIR_PATH, "hosts")):
- try:
- open(path_join_robust(BASEDIR_PATH, "hosts"), "w+").close()
- except Exception:
- print_failure("ERROR: No 'hosts' file in the folder,"
- "try creating one manually")
+ # Create a hosts file if it doesn't exist.
+ hosts_file = path_join_robust(BASEDIR_PATH, "hosts")
- if not settings["freshen"]:
+ if not os.path.isfile(hosts_file):
+ try:
+ open(hosts_file, "w+").close()
+ except (IOError, OSError):
+ # Starting in Python 3.3, IOError is aliased
+ # OSError. However, we have to catch both for
+ # Python 2.x failures.
+ print_failure("ERROR: No 'hosts' file in the folder. "
+ "Try creating one manually.")
+
+ if not freshen:
return
prompt = "Do you want to update all data sources?"
- if settings["auto"] or query_yes_no(prompt):
+
+ if update_auto or query_yes_no(prompt):
update_all_sources()
- elif not settings["auto"]:
+ elif not update_auto:
print("OK, we'll stick with what we've got locally.")
-def prompt_for_exclusions():
+def prompt_for_exclusions(skip_prompt):
"""
Prompt the user to exclude any custom domains from being blocked.
+
+ Parameters
+ ----------
+ skip_prompt : bool
+ Whether or not to skip prompting for custom domains to be excluded.
+ If true, the function returns immediately.
"""
prompt = ("Do you want to exclude any domains?\n"
"For example, hulu.com video streaming must be able to access "
"its tracking and ad servers in order to play video.")
- if not settings["auto"]:
+ if not skip_prompt:
if query_yes_no(prompt):
display_exclusion_options()
else:
print("OK, we'll only exclude domains in the whitelist.")
-def prompt_for_flush_dns_cache():
+def prompt_for_flush_dns_cache(flush_cache, prompt_flush):
"""
Prompt the user to flush the DNS cache.
+
+ Parameters
+ ----------
+ flush_cache : bool
+ Whether to flush the DNS cache without prompting.
+ prompt_flush : bool
+ If `flush_cache` is False, whether we should prompt for flushing the
+ cache. Otherwise, the function returns immediately.
"""
- if settings["flushdnscache"]:
+ if flush_cache:
flush_dns_cache()
-
- if not settings["auto"]:
+ elif prompt_flush:
if query_yes_no("Attempt to flush the DNS cache?"):
flush_dns_cache()
-def prompt_for_move(final_file):
+def prompt_for_move(final_file, **move_params):
"""
Prompt the user to move the newly created hosts file to its designated
location in the OS.
----------
final_file : file
The file object that contains the newly created hosts data.
+ move_params : kwargs
+ Dictionary providing additional parameters for moving the hosts file
+ into place. Currently, those fields are:
+
+ 1) auto
+ 2) replace
+ 3) skipstatichosts
+
+ Returns
+ -------
+ move_file : bool
+ Whether or not the final hosts file was moved.
"""
- if settings["replace"] and not settings["skipstatichosts"]:
+ skip_static_hosts = move_params["skipstatichosts"]
+
+ if move_params["replace"] and not skip_static_hosts:
move_file = True
- elif settings["auto"] or settings["skipstatichosts"]:
+ elif move_params["auto"] or skip_static_hosts:
move_file = False
else:
prompt = ("Do you want to replace your existing hosts file " +
if move_file:
move_hosts_file_into_place(final_file)
- prompt_for_flush_dns_cache()
+
+ return move_file
# End Prompt the User