class HostManipulationTestCase(unittest.TestCase): def setUp(self): fh, self.hosts_file = tempfile.mkstemp() self.hosts = Hosts(self.hosts_file) def tearDown(self): os.unlink(self.hosts_file) def assertHasHostLine(self, host_line): if not has_host_line(self.hosts_file, host_line): raise AssertionError("Line not present: {0}".format(host_line)) def assertDoesNotHaveHostLine(self, host_line): if has_host_line(self.hosts_file, host_line): raise AssertionError("Line is present: {0}".format(host_line)) def test_set_one(self): self.hosts.set_one("test", "1.2.3.4") self.hosts.write(self.hosts_file) self.assertHasHostLine("1.2.3.4 test") def test_set_all(self): self.hosts.set_all(["test", "alias"], "1.2.3.4") self.hosts.write(self.hosts_file) self.assertHasHostLine("1.2.3.4 test") self.assertHasHostLine("1.2.3.4 alias") def test_remove_one(self): self.hosts.set_one("test", "1.2.3.4") self.hosts.write(self.hosts_file) self.hosts = Hosts(self.hosts_file) self.hosts.remove_one("test") self.hosts.write(self.hosts_file) self.assertDoesNotHaveHostLine("1.2.3.4 test") def test_remove_no_raise(self): self.assertRaises(KeyError, self.hosts.remove_one, "test") self.hosts.remove_one("test", False) def test_purge_empty_records(self): with open(self.hosts_file, 'a') as f: f.write('1.2.3.4') self.assertHasHostLine('1.2.3.4') hosts = Hosts(self.hosts_file) hosts.set_one("test", "1.2.3.4") hosts.set_one("", "1.2.3.4") hosts.write(self.hosts_file) self.assertDoesNotHaveHostLine('1.2.3.4')
def test_purge_empty_records(self): with open(self.hosts_file, 'a') as f: f.write('1.2.3.4') self.assertHasHostLine('1.2.3.4') hosts = Hosts(self.hosts_file) hosts.set_one("test", "1.2.3.4") hosts.set_one("", "1.2.3.4") hosts.write(self.hosts_file) self.assertDoesNotHaveHostLine('1.2.3.4')
#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ __author__ = 'wenin819' """ from hosts import Hosts def_hosts = { "test": "1.2.3.6" } if __name__ == '__main__': import os import argparse if os.name == 'nt': hosts_path = os.path.join(os.environ['SYSTEMROOT'], 'system32/drivers/etc/hosts') elif os.name == 'posix': hosts_path = '/etc/hosts' else: raise Exception('Unsupported OS: %s' % os.name) hosts = Hosts(hosts_path) for (host, val) in def_hosts.items(): hosts.set_one(host, val) hosts.write(hosts_path)