oneforall-kjl 0.1.1__py3-none-any.whl

OneForAll/common/database.py

@@ -0,0 +1,277 @@
+"""
+SQLite database initialization and operation
+"""
+
+from common import records
+
+from common.records import Connection
+from config.log import logger
+from config import settings
+
+
+class Database(object):
+    def __init__(self, db_path=None):
+        self.conn = self.get_conn(db_path)
+
+    @staticmethod
+    def get_conn(db_path):
+        """
+        Get database connection
+
+        :param   db_path: Database path
+        :return: db_conn: SQLite database connection
+        """
+        logger.log('TRACE', f'Establishing database connection')
+        if isinstance(db_path, Connection):
+            return db_path
+        protocol = 'sqlite:///'
+        if not db_path:  # 数据库路径为空连接默认数据库
+            db_path = f'{protocol}{settings.result_save_dir}/result.sqlite3'
+        else:
+            db_path = f'{protocol}{db_path}'
+        db = records.Database(db_path)  # 不存在数据库时会新建一个数据库
+        logger.log('TRACE', f'Use the database: {db_path}')
+        return db.get_connection()
+
+    def query(self, sql):
+        try:
+            results = self.conn.query(sql)
+        except Exception as e:
+            logger.log('ERROR', e.args)
+            return None
+        return results
+
+    def create_table(self, table_name):
+        """
+        Create table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        if self.exist_table(table_name):
+            logger.log('TRACE', f'{table_name} table already exists')
+            return
+        logger.log('TRACE', f'Creating {table_name} table')
+        self.query(f'create table "{table_name}" ('
+                   f'id integer primary key,'
+                   f'alive int,'
+                   f'request int,'
+                   f'resolve int,'
+                   f'url text,'
+                   f'subdomain text,'
+                   f'port int,'
+                   f'level int,'
+                   f'cname text,'
+                   f'ip text,'
+                   f'public int,'
+                   f'cdn int,'
+                   f'status int,'
+                   f'reason text,'
+                   f'title text,'
+                   f'banner text,'
+                   f'header text,'
+                   f'history text,'
+                   f'response text,'
+                   f'ip_times text,'
+                   f'cname_times text,'
+                   f'ttl text,'
+                   f'cidr text,'
+                   f'asn text,'
+                   f'org text,'
+                   f'addr text,'
+                   f'isp text,'
+                   f'resolver text,'
+                   f'module text,'
+                   f'source text,'
+                   f'elapse float,'
+                   f'find int)')
+
+    def insert_table(self, table_name, result):
+        table_name = table_name.replace('.', '_')
+        self.conn.query(
+            f'insert into "{table_name}" '
+            f'(id, alive, resolve, request, url, subdomain, port, level,'
+            f'cname, ip, public, cdn, status, reason, title, banner, header,'
+            f'history, response, ip_times, cname_times, ttl, cidr, asn, org,'
+            f'addr, isp, resolver, module, source, elapse, find) '
+            f'values (:id, :alive, :resolve, :request, :url,'
+            f':subdomain, :port, :level, :cname, :ip, :public, :cdn,'
+            f':status, :reason, :title, :banner, :header, :history, :response,'
+            f':ip_times, :cname_times, :ttl, :cidr, :asn, :org, :addr, :isp,'
+            f':resolver, :module, :source, :elapse, :find)', **result)
+
+    def save_db(self, table_name, results, module_name=None):
+        """
+        Save the results of each module in the database
+
+        :param str table_name: table name
+        :param list results: results list
+        :param str module_name: module
+        """
+        logger.log('TRACE', f'Saving the subdomain results of {table_name} '
+                            f'found by module {module_name} into database')
+        table_name = table_name.replace('.', '_')
+        if results:
+            try:
+                self.conn.bulk_query(
+                    f'insert into "{table_name}" '
+                    f'(id, alive, resolve, request, url, subdomain, port, level, '
+                    f'cname, ip, public, cdn, status, reason, title, banner, header, '
+                    f'history, response, ip_times, cname_times, ttl, cidr, asn, org, '
+                    f'addr, isp, resolver, module, source, elapse, find) '
+                    f'values (:id, :alive, :resolve, :request, :url, '
+                    f':subdomain, :port, :level, :cname, :ip, :public, :cdn,'
+                    f':status, :reason, :title, :banner, :header, :history, :response, '
+                    f':ip_times, :cname_times, :ttl, :cidr, :asn, :org, :addr, :isp, '
+                    f':resolver, :module, :source, :elapse, :find)', results)
+            except Exception as e:
+                logger.log('ERROR', e)
+
+    def exist_table(self, table_name):
+        """
+        Determine table exists
+
+        :param   str table_name: table name
+        :return  bool: Whether table exists
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Determining whether the {table_name} table exists')
+        results = self.query(f'select count() from sqlite_master where type = "table" and'
+                             f' name = "{table_name}"')
+        if results.scalar() == 0:
+            return False
+        else:
+            return True
+
+    def copy_table(self, table_name, bak_table_name):
+        """
+        Copy table to create backup
+
+        :param str table_name: table name
+        :param str bak_table_name: new table name
+        """
+        table_name = table_name.replace('.', '_')
+        bak_table_name = bak_table_name.replace('.', '_')
+        logger.log('TRACE', f'Copying {table_name} table to {bak_table_name} new table')
+        self.query(f'drop table if exists "{bak_table_name}"')
+        self.query(f'create table "{bak_table_name}" '
+                   f'as select * from "{table_name}"')
+
+    def clear_table(self, table_name):
+        """
+        Clear the table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Clearing data in table {table_name}')
+        self.query(f'delete from "{table_name}"')
+
+    def drop_table(self, table_name):
+        """
+        Delete table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Deleting {table_name} table')
+        self.query(f'drop table if exists "{table_name}"')
+
+    def rename_table(self, table_name, new_table_name):
+        """
+        Rename table name
+
+        :param str table_name: old table name
+        :param str new_table_name: new table name
+        """
+        table_name = table_name.replace('.', '_')
+        new_table_name = new_table_name.replace('.', '_')
+        logger.log('TRACE', f'Renaming {table_name} table to {new_table_name} table')
+        self.query(f'alter table "{table_name}" '
+                   f'rename to "{new_table_name}"')
+
+    def deduplicate_subdomain(self, table_name):
+        """
+        Deduplicate subdomains in the table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Deduplicating subdomains in {table_name} table')
+        self.query(f'delete from "{table_name}" where '
+                   f'id not in (select min(id) '
+                   f'from "{table_name}" group by subdomain)')
+
+    def remove_invalid(self, table_name):
+        """
+        Remove nulls or invalid subdomains in the table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Removing invalid subdomains in {table_name} table')
+        self.query(f'delete from "{table_name}" where '
+                   f'subdomain is null or resolve == 0')
+
+    def get_data(self, table_name):
+        """
+        Get all the data in the table
+
+        :param str table_name: table name
+        """
+        table_name = table_name.replace('.', '_')
+        logger.log('TRACE', f'Get all the data from {table_name} table')
+        return self.query(f'select * from "{table_name}"')
+
+    def export_data(self, table_name, alive, limit):
+        """
+        Get part of the data in the table
+
+        :param str table_name: table name
+        :param any alive: alive flag
+        :param str limit: limit value
+        """
+        table_name = table_name.replace('.', '_')
+        sql = f'select id, alive, request, resolve, url, subdomain, level,' \
+              f'cname, ip, public, cdn, port, status, reason, title, banner,' \
+              f'cidr, asn, org, addr, isp, source from "{table_name}" '
+        if alive and limit:
+            if limit in ['resolve', 'request']:
+                where = f' where {limit} = 1'
+                sql += where
+        elif alive:
+            where = f' where alive = 1'
+            sql += where
+        sql += ' order by subdomain'
+        logger.log('TRACE', f'Get the data from {table_name} table')
+        return self.query(sql)
+
+    def count_alive(self, table_name):
+        table_name = table_name.replace('.', '_')
+        sql = f'select count() from "{table_name}" where alive = 1'
+        return self.query(sql)
+
+    def get_resp_by_url(self, table_name, url):
+        table_name = table_name.replace('.', '_')
+        sql = f'select response from "{table_name}" where url = "{url}"'
+        logger.log('TRACE', f'Get response data from {url}')
+        return self.query(sql).scalar()
+
+    def get_data_by_fields(self, table_name, fields):
+        table_name = table_name.replace('.', '_')
+        field_str = ', '.join(fields)
+        sql = f'select {field_str} from "{table_name}"'
+        logger.log('TRACE', f'Get specified field data {fields} from {table_name} table')
+        return self.query(sql)
+
+    def update_data_by_url(self, table_name, info, url):
+        table_name = table_name.replace('.', '_')
+        field_str = ', '.join(map(lambda kv: f'{kv[0]} = "{kv[1]}"', info.items()))
+        sql = f'update "{table_name}" set {field_str} where url = "{url}"'
+        return self.query(sql)
+
+    def close(self):
+        """
+        Close the database connection
+        """
+        self.conn.close()

OneForAll/common/domain.py

@@ -0,0 +1,63 @@
+import re
+from common import tldextract
+from config import settings
+
+
+class Domain(object):
+    """
+    Processing domain class
+
+    :param str string: input string
+    """
+    def __init__(self, string):
+        self.string = str(string)
+        self.regexp = r'\b((?=[a-z0-9-]{1,63}\.)(xn--)?[a-z0-9]+(-[a-z0-9]+)*\.)+[a-z]{2,63}\b'
+        self.domain = None
+
+    def match(self):
+        """
+        match domain
+
+        :return : result
+        """
+        result = re.search(self.regexp, self.string, re.I)
+        if result:
+            return result.group()
+        return None
+
+    def extract(self):
+        """
+        extract domain
+
+        >>> d = Domain('www.example.com')
+        <domain.Domain object>
+        >>> d.extract()
+        ExtractResult(subdomain='www', domain='example', suffix='com')
+
+        :return: extracted domain results
+        """
+        data_storage_dir = settings.data_storage_dir
+        extract_cache_file = data_storage_dir.joinpath('public_suffix_list.dat')
+        ext = tldextract.TLDExtract(extract_cache_file)
+        result = self.match()
+        if result:
+            return ext(result)
+        return None
+
+    def registered(self):
+        """
+        registered domain
+
+        >>> d = Domain('www.example.com')
+        <domain.Domain object>
+        >>> d.registered()
+        example.com
+
+        :return: registered domain result
+        """
+        if not settings.use_tld_extract:
+            return self.string
+        result = self.extract()
+        if result:
+            return result.registered_domain
+        return None

OneForAll/common/ipasn.py

@@ -0,0 +1,42 @@
+import zipfile
+
+from common.utils import ip_to_int
+from config.setting import data_storage_dir
+from common.database import Database
+
+
+def get_db_path():
+    zip_path = data_storage_dir.joinpath('ip2location.zip')
+    db_path = data_storage_dir.joinpath('ip2location.db')
+    if db_path.exists():
+        return db_path
+    zf = zipfile.ZipFile(str(zip_path))
+    zf.extract('ip2location.db', data_storage_dir)
+    return db_path
+
+
+class IPAsnInfo(Database):
+    def __init__(self):
+        path = get_db_path()
+        Database.__init__(self, path)
+
+    def find(self, ip):
+        info = {'cidr': '', 'asn': '', 'org': ''}
+        if isinstance(ip, (int, str)):
+            ip = ip_to_int(ip)
+        else:
+            return info
+        sql = f'SELECT * FROM asn WHERE ip_from <= {ip} AND ip_to >= {ip} LIMIT 1;'
+        result = self.query(sql)
+        if not hasattr(result, 'dataset'):
+            return info
+        asn = result.as_dict()
+        info['cidr'] = asn[0]['cidr']
+        info['asn'] = f"AS{asn[0]['asn']}"
+        info['org'] = asn[0]['as']
+        return info
+
+
+if __name__ == "__main__":
+    asn_info = IPAsnInfo()
+    print(asn_info.find("188.81.94.77"))

OneForAll/common/ipreg.py

@@ -0,0 +1,139 @@
+"""
+" ip2region python searcher client module
+"
+" Author: koma<komazhang@foxmail.com>
+" Date : 2015-11-06
+"""
+import io
+import sys
+import socket
+import struct
+
+from config import settings
+
+
+class IpRegInfo(object):
+    __INDEX_BLOCK_LENGTH = 12
+    __TOTAL_HEADER_LENGTH = 8192
+
+    __f = None
+    __headerSip = []
+    __headerPtr = []
+    __headerLen = 0
+    __indexSPtr = 0
+    __indexLPtr = 0
+    __indexCount = 0
+    __dbBinStr = ''
+
+    def __init__(self, db_file):
+        self.init_database(db_file)
+
+    def memory_search(self, ip):
+        """
+        " memory search method
+        " param: ip
+        """
+        if not ip.isdigit():
+            ip = self.ip2long(ip)
+
+        if self.__dbBinStr == '':
+            self.__dbBinStr = self.__f.read()  # read all the contents in file
+            self.__indexSPtr = self.get_long(self.__dbBinStr, 0)
+            self.__indexLPtr = self.get_long(self.__dbBinStr, 4)
+            self.__indexCount = int((self.__indexLPtr - self.__indexSPtr) /
+                                    self.__INDEX_BLOCK_LENGTH) + 1
+
+        l, h, data_ptr = (0, self.__indexCount, 0)
+        while l <= h:
+            m = int((l + h) >> 1)
+            p = self.__indexSPtr + m * self.__INDEX_BLOCK_LENGTH
+            sip = self.get_long(self.__dbBinStr, p)
+
+            if ip < sip:
+                h = m - 1
+            else:
+                eip = self.get_long(self.__dbBinStr, p + 4)
+                if ip > eip:
+                    l = m + 1
+                else:
+                    data_ptr = self.get_long(self.__dbBinStr, p + 8)
+                    break
+
+        if data_ptr == 0:
+            raise Exception("Data pointer not found")
+
+        return self.return_data(data_ptr)
+
+    def init_database(self, db_file):
+        """
+        " initialize the database for search
+        " param: dbFile
+        """
+        try:
+            self.__f = io.open(db_file, "rb")
+        except IOError as e:
+            print("[Error]: %s" % e)
+            sys.exit()
+
+    def return_data(self, data_ptr):
+        """
+        " get ip data from db file by data start ptr
+        " param: data ptr
+        """
+        data_len = (data_ptr >> 24) & 0xFF
+        data_ptr = data_ptr & 0x00FFFFFF
+
+        self.__f.seek(data_ptr)
+        data = self.__f.read(data_len)
+
+        info = {"city_id": self.get_long(data, 0),
+                "region": data[4:].decode('utf-8')}
+        return info
+
+    @staticmethod
+    def ip2long(ip):
+        _ip = socket.inet_aton(ip)
+        return struct.unpack("!L", _ip)[0]
+
+    @staticmethod
+    def is_ip(ip):
+        p = ip.split(".")
+        if len(p) != 4:
+            return False
+        for pp in p:
+            if not pp.isdigit():
+                return False
+            if len(pp) > 3:
+                return False
+            if int(pp) > 255:
+                return False
+        return True
+
+    @staticmethod
+    def get_long(b, offset):
+        if len(b[offset:offset + 4]) == 4:
+            return struct.unpack('I', b[offset:offset + 4])[0]
+        return 0
+
+    def close(self):
+        if self.__f is not None:
+            self.__f.close()
+        self.__dbBinStr = None
+        self.__headerPtr = None
+        self.__headerSip = None
+
+
+class IpRegData(IpRegInfo):
+    def __init__(self):
+        path = settings.data_storage_dir.joinpath('ip2region.db')
+        IpRegInfo.__init__(self, path)
+
+    def query(self, ip):
+        result = self.memory_search(ip)
+        addr_list = result.get('region').split('|')
+        addr = ''.join(filter(lambda x: x != '0', addr_list[:-1]))
+        isp = addr_list[-1]
+        if isp == '0':
+            isp = '未知'
+        info = {'addr': addr, 'isp': isp}
+        return info

OneForAll/common/lookup.py

@@ -0,0 +1,28 @@
+from common.module import Module
+from common import utils
+from config.log import logger
+
+
+class Lookup(Module):
+    """
+    DNS query base class
+    """
+
+    def __init__(self):
+        Module.__init__(self)
+        self.qtype = ''
+
+    def query(self):
+        """
+        Query the TXT record of domain
+        :return: query result
+        """
+        answer = utils.dns_query(self.domain, self.qtype)
+        if answer is None:
+            return None
+        for item in answer:
+            record = item.to_text()
+            subdomains = self.match_subdomains(record)
+            self.subdomains.update(subdomains)
+            logger.log('DEBUG', record)
+        return self.subdomains