From 515eb5756cd483a8fe230c6423652e0cbc793931 Mon Sep 17 00:00:00 2001 From: Nikita Petrov Date: Sun, 22 Apr 2018 01:48:23 +0300 Subject: [PATCH] Code refactor and performance improvements - The checker fully refactored. Now it could be used as external lib in Python projects - Sqlite now used as cache backed. This significantly reduced the run duration of checks --- .gitignore | 3 +- check.py | 111 ------------------ checker.py | 317 ++++++++++++++++++++++++++++++++++++++++++++++++++ entrypoint.sh | 6 +- 4 files changed, 322 insertions(+), 115 deletions(-) delete mode 100755 check.py create mode 100755 checker.py diff --git a/.gitignore b/.gitignore index 0471460..a7a0e5d 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ venv *.pyc -registry.zip +cache +.idea diff --git a/check.py b/check.py deleted file mode 100755 index 3c4dc9b..0000000 --- a/check.py +++ /dev/null @@ -1,111 +0,0 @@ -#!/usr/bin/env python3 - -import argparse -import os -import shutil -import urllib.request -import zipfile -import json -import socket - -import netaddr - - -arg_parser = argparse.ArgumentParser(description="Roskomnadzor ban registry checker") - -arg_parser.add_argument("mode", type=str, choices=("fetch", "check"), help="run mode") -arg_parser.add_argument("host", type=str, nargs="*", - help="ip address, network or fqdn to check") -arg_parser.add_argument("--registry-url", type=str, - default="https://github.com/zapret-info/z-i/archive/master.zip", - help="Registry repository url") - -script_args = arg_parser.parse_args() - -if script_args.mode == "check" and not script_args.host: - arg_parser.error("the host argument is required for check mode") - - -script_dir = os.path.dirname(__file__) -registry_dir = os.path.join(script_dir, "registry") -registry_archive = os.path.join(registry_dir, "registry.zip") -ips_data_file = os.path.join(registry_dir, "z-i-master/dump.csv") -fqdn_data_file = os.path.join(registry_dir, "z-i-master/nxdomain.txt") - - -def fetch(): - if not os.path.exists(registry_dir): - os.mkdir(registry_dir) - - for file in os.listdir(registry_dir): - file_path = os.path.join(registry_dir, file) - - if os.path.isfile(file_path): - os.remove(file_path) - else: - shutil.rmtree(file_path) - - urllib.request.urlretrieve(script_args.registry_url, registry_archive) - - zip_f = zipfile.ZipFile(registry_archive, 'r') - zip_f.extractall(registry_dir) - zip_f.close() - - - -def check(hosts): - result = {} - - ips_data = open(ips_data_file, encoding="cp1251").readlines()[1:] - fqdn_data = open(fqdn_data_file).readlines() - - for host in hosts: - try: - host_obj = netaddr.IPNetwork(host) - except netaddr.core.AddrFormatError: - host_obj = host - - ip_objs = [] - - if type(host_obj) is str: - try: - host_ips = socket.gethostbyname_ex(host_obj)[2] - - for host_ip in host_ips: - ip_objs.append(netaddr.IPNetwork(host_ip)) - except (ValueError, socket.gaierror): - host_ip = None - - for line in fqdn_data: - fqdn_reg = line.strip() - - if fqdn_reg == host_obj: - result[host] = result.get(host, []) - result[host].append(fqdn_reg) - else: - ip_objs = [host_obj] - - for ip_obj in ip_objs: - for line in ips_data: - subnets = line.strip().split(";")[0].split("|") - - for subnet in subnets: - subnet_reg = subnet.strip() - - try: - subnet_obj = netaddr.IPNetwork(subnet_reg) - except netaddr.core.AddrFormatError: - continue - - if ip_obj in subnet_obj: - result[host] = result.get(host, []) - result[host].append(subnet_reg) - - if result: - print(json.dumps(result)) - - -if script_args.mode == "fetch": - fetch() -elif script_args.mode == "check": - check(script_args.host) diff --git a/checker.py b/checker.py new file mode 100755 index 0000000..df62972 --- /dev/null +++ b/checker.py @@ -0,0 +1,317 @@ +#!/usr/bin/env python3 + +import argparse +import os +import shutil +import urllib.request +import zipfile +import json +import socket +import glob +import sqlite3 +import math + +import netaddr + + +class RknChecker(object): + def __init__(self, cache_dir=None): + app_dir = os.path.dirname(__file__) + + self.cache_dir = cache_dir or os.path.join(app_dir, "cache") + self.registry_db_file = os.path.join(self.cache_dir, "registry.db") + + def fetch(self, registry_url): + registry_dir = os.path.join(self.cache_dir, "registry") + registry_db_file_temp = "{}.tmp".format(self.registry_db_file) + + self._mkdir_if_not_exists(self.cache_dir) + + self._fetch_registry(registry_url, registry_dir) + self._fill_database(registry_dir, registry_db_file_temp) + + self._commit_fetch(registry_db_file_temp) + + @staticmethod + def _mkdir_if_not_exists(dir_path): + if not os.path.exists(dir_path): + os.mkdir(dir_path) + + @staticmethod + def _rm_file_if_exists(file_path): + if os.path.exists(file_path): + os.remove(file_path) + + def _fetch_registry(self, registry_url, registry_dir): + registry_arch_file = os.path.join(self.cache_dir, "registry.zip") + + urllib.request.urlretrieve(registry_url, registry_arch_file) + + self._mkdir_if_not_exists(registry_dir) + self._unzip_file(registry_arch_file, registry_dir) + + @staticmethod + def _unzip_file(file_path, dir_path): + zip_f = zipfile.ZipFile(file_path, "r") + zip_f.extractall(dir_path) + zip_f.close() + + def _fill_database(self, registry_dir, db_file_path): + self._rm_file_if_exists(db_file_path) + db_conn = sqlite3.connect(db_file_path) + + self._prepare_database(db_conn) + + self._fill_db_ips_data(db_conn, registry_dir) + self._fill_db_fqdn_data(db_conn, registry_dir) + + db_conn.close() + + @staticmethod + def _prepare_database(db_conn): + db_cursor = db_conn.cursor() + + db_cursor.execute("CREATE TABLE ip_networks (start_addr INTEGER, end_addr INTEGER)") + db_cursor.execute("CREATE INDEX ip_netw_start_addr ON ip_networks (start_addr)") + db_cursor.execute("CREATE INDEX ip_netw_end_addr ON ip_networks (end_addr)") + + db_cursor.execute("CREATE TABLE fqdns (fqdn TEXT)") + db_cursor.execute("CREATE INDEX fqdns_fqdn ON fqdns (fqdn)") + + def _fill_db_ips_data(self, db_conn, registry_dir): + registry_ips_file = glob.glob(os.path.join(registry_dir, "**", "dump.csv"), recursive=True)[0] + + ips_data = self._load_ips_data(registry_ips_file) + self._save_db_data(ips_data, db_conn, "INSERT INTO ip_networks (start_addr, end_addr) VALUES (?, ?)", + lambda r: (r["start_addr"], r["end_addr"])) + + @staticmethod + def _load_ips_data(file_path): + data = [] + + file_lines = open(file_path, encoding="cp1251").readlines()[1:] + + for line in file_lines: + networks = line.strip().split(";")[0].split("|") + + for network in networks: + network_reg = network.strip() + + try: + network_obj = netaddr.IPNetwork(network_reg) + except netaddr.core.AddrFormatError: + continue + + data_rec = { + "start_addr": network_obj.first, + "end_addr": network_obj.last, + } + + data.append(data_rec) + + return data + + def _fill_db_fqdn_data(self, db_conn, registry_dir): + registry_fqdn_file = glob.glob(os.path.join(registry_dir, "**", "nxdomain.txt"), recursive=True)[0] + + fqdn_data = self._load_fqdn_data(registry_fqdn_file) + self._save_db_data(fqdn_data, db_conn, "INSERT INTO fqdns (fqdn) VALUES (?)", lambda r: (r,)) + + @staticmethod + def _save_db_data(data, db_conn, sql, parameters_fn=None): + db_cursor = db_conn.cursor() + + for data_rec in data: + if parameters_fn: + parameters = parameters_fn(data_rec) + else: + parameters = data_rec + + db_cursor.execute(sql, parameters) + + db_conn.commit() + + @staticmethod + def _load_fqdn_data(file_path): + data = [] + + file_lines = open(file_path).readlines() + + for line in file_lines: + fqdn_reg = line.strip().lower() + data.append(fqdn_reg) + + return data + + def _commit_fetch(self, registry_db_file_temp): + os.rename(registry_db_file_temp, self.registry_db_file) + + self._cleanup_cache_dir() + + def _cleanup_cache_dir(self): + for node in os.listdir(self.cache_dir): + node_path = os.path.join(self.cache_dir, node) + + if os.path.isfile(node_path): + if node_path == self.registry_db_file: + continue + + os.remove(node_path) + else: + shutil.rmtree(node_path) + + def check(self, host): + hosts = host if hasattr(host, "__iter__") else [host] + db_conn = sqlite3.connect(self.registry_db_file) + + result = self._check_hosts(hosts, db_conn) + db_conn.close() + + return result + + def _check_hosts(self, hosts, db_conn): + result = {} + + for host in hosts: + host_results = self._check_host(host, db_conn) + + if host_results: + result[host] = host_results + + return result + + def _check_host(self, host, db_conn): + results = [] + + host_objs = self._get_host_objs(host) + + for host_obj in host_objs: + host_obj_results = self._check_host_obj(host_obj, db_conn) + + if host_obj_results: + results += host_obj_results + + return results + + def _get_host_objs(self, host): + host_objs = [] + + host_obj = self._get_host_obj(host) + host_objs.append(host_obj) + + if isinstance(host_obj, str): + ip_objs = self._get_fqdn_ip_objs(host_obj) + host_objs += ip_objs + + return host_objs + + @staticmethod + def _get_host_obj(host): + try: + host_obj = netaddr.IPAddress(host) + except netaddr.core.AddrFormatError: + host_obj = host + + return host_obj + + @staticmethod + def _get_fqdn_ip_objs(fqdn): + ip_objs = [] + + try: + fqdn_ips = socket.gethostbyname_ex(fqdn)[2] + + for fqdn_ip in fqdn_ips: + ip_obj = netaddr.IPAddress(fqdn_ip) + ip_objs.append(ip_obj) + except (ValueError, socket.gaierror): + pass + + return ip_objs + + def _check_host_obj(self, host_obj, db_conn): + results = [] + + if isinstance(host_obj, str): + result = self._check_fqdn(host_obj, db_conn) + + if result: + results = [result] + else: + results = self._check_ip_obj(host_obj, db_conn) + + return results + + @staticmethod + def _check_fqdn(fqdn, db_conn): + result = None + db_cursor = db_conn.cursor() + + db_cursor.execute("SELECT fqdn FROM fqdns WHERE fqdn = ?", (fqdn.lower(),)) + select_result = db_cursor.fetchone() + + if select_result: + result = select_result[0] + + return result + + def _check_ip_obj(self, ip_obj, db_conn): + results = [] + db_cursor = db_conn.cursor() + + db_cursor.execute("SELECT start_addr, end_addr FROM ip_networks WHERE ? BETWEEN start_addr AND end_addr", + (int(ip_obj),)) + select_results = db_cursor.fetchall() + + for select_result in select_results: + start_addr = select_result[0] + end_addr = select_result[1] + + network_obj = self._get_network_obj(start_addr, end_addr) + + results.append(str(network_obj)) + + return results + + @staticmethod + def _get_network_obj(start_addr, end_addr): + network_addr = start_addr + network_mask = 32 - round(math.log2(end_addr - start_addr + 1)) + + network_obj = netaddr.IPNetwork((network_addr, network_mask)) + + return network_obj + + +def parse_args(): + arg_parser = argparse.ArgumentParser(description="Roskomnadzor prohibited resources registry checker") + + arg_parser.add_argument("mode", type=str, choices=("fetch", "check"), help="run mode") + arg_parser.add_argument("host", type=str, nargs="*", + help="ip address, network or fqdn to check") + arg_parser.add_argument("--registry-url", type=str, + default="https://github.com/zapret-info/z-i/archive/master.zip", + help="registry url") + + args = arg_parser.parse_args() + + if args.mode == "check" and not args.host: + arg_parser.error("at least one host argument is required in check mode") + + return args + + +if __name__ == "__main__": + args = parse_args() + result = None + + rkn_checker = RknChecker() + + if args.mode == "fetch": + rkn_checker.fetch(args.registry_url) + elif args.mode == "check": + result = rkn_checker.check(args.host) + + if result: + json_data = json.dumps(result) + print(json_data) \ No newline at end of file diff --git a/entrypoint.sh b/entrypoint.sh index 038efe7..e7b83af 100755 --- a/entrypoint.sh +++ b/entrypoint.sh @@ -1,8 +1,8 @@ #!/bin/sh if [ ! -z "${REGISTRY_URL}" ]; then - python check.py --registry-url=${REGISTRY_URL} fetch + python checker.py --registry-url=${REGISTRY_URL} fetch else - python check.py fetch + python checker.py fetch fi -python check.py check "$@" +python checker.py check "$@"