oneforall-kjl 0.1.1__py3-none-any.whl

OneForAll/common/records.py

@@ -0,0 +1,363 @@
+import os
+from collections import OrderedDict
+from inspect import isclass
+
+from sqlalchemy import create_engine, exc, inspect, text
+
+from .tablib import tablib
+
+DATABASE_URL = os.environ.get('DATABASE_URL')
+
+
+def is_exception(obj):
+    """Given an object, return a boolean indicating whether it is an instance
+    or subclass of :py:class:`Exception`.
+    """
+    if isinstance(obj, Exception):
+        return True
+    if isclass(obj) and issubclass(obj, Exception):
+        return True
+    return False
+
+
+class Record(object):
+    """A row, from a query, from a database."""
+    __slots__ = ('_keys', '_values')
+
+    def __init__(self, keys, values):
+        self._keys = keys
+        self._values = values
+
+        # Ensure that lengths match properly.
+        assert len(self._keys) == len(self._values)
+
+    def keys(self):
+        """Returns the list of column names from the query."""
+        return self._keys
+
+    def values(self):
+        """Returns the list of values from the query."""
+        return self._values
+
+    def __repr__(self):
+        return '<Record {}>'.format(self.export('json')[1:-1])
+
+    def __getitem__(self, key):
+        # Support for index-based lookup.
+        if isinstance(key, int):
+            return self.values()[key]
+
+        # Support for string-based lookup.
+        if key in self.keys():
+            i = self.keys().index(key)
+            if self.keys().count(key) > 1:
+                raise KeyError("Record contains multiple '{}' fields.".format(key))
+            return self.values()[i]
+
+        raise KeyError("Record contains no '{}' field.".format(key))
+
+    def __getattr__(self, key):
+        try:
+            return self[key]
+        except KeyError as e:
+            raise AttributeError(e)
+
+    def __dir__(self):
+        standard = dir(super(Record, self))
+        # Merge standard attrs with generated ones (from column names).
+        return sorted(standard + [str(k) for k in self.keys()])
+
+    def get(self, key, default=None):
+        """Returns the value for a given key, or default."""
+        try:
+            return self[key]
+        except KeyError:
+            return default
+
+    def as_dict(self, ordered=False):
+        """Returns the row as a dictionary, as ordered."""
+        items = zip(self.keys(), self.values())
+
+        return OrderedDict(items) if ordered else dict(items)
+
+    @property
+    def dataset(self):
+        """A Tablib Dataset containing the row."""
+        data = tablib.Dataset()
+        data.headers = self.keys()
+
+        row = _reduce_datetimes(self.values())
+        data.append(row)
+
+        return data
+
+    def export(self, format, **kwargs):
+        """Exports the row to the given format."""
+        return self.dataset.export(format, **kwargs)
+
+
+class RecordCollection(object):
+    """A set of excellent Records from a query."""
+
+    def __init__(self, rows):
+        self._rows = rows
+        self._all_rows = []
+        self.pending = True
+
+    def __repr__(self):
+        return '<RecordCollection size={} pending={}>'.format(len(self), self.pending)
+
+    def __iter__(self):
+        """Iterate over all rows, consuming the underlying generator
+        only when necessary."""
+        i = 0
+        while True:
+            # Other code may have iterated between yields,
+            # so always check the cache.
+            if i < len(self):
+                yield self[i]
+            else:
+                # Throws StopIteration when done.
+                # Prevent StopIteration bubbling from generator,
+                # following https://www.python.org/dev/peps/pep-0479/
+                try:
+                    yield next(self)
+                except StopIteration:
+                    return
+            i += 1
+
+    def next(self):
+        return self.__next__()
+
+    def __next__(self):
+        try:
+            nextrow = next(self._rows)
+            self._all_rows.append(nextrow)
+            return nextrow
+        except StopIteration:
+            self.pending = False
+            raise StopIteration('RecordCollection contains no more rows.')
+
+    def __getitem__(self, key):
+        is_int = isinstance(key, int)
+
+        # Convert RecordCollection[1] into slice.
+        if is_int:
+            key = slice(key, key + 1)
+
+        while len(self) < key.stop or key.stop is None:
+            try:
+                next(self)
+            except StopIteration:
+                break
+
+        rows = self._all_rows[key]
+        if is_int:
+            return rows[0]
+        else:
+            return RecordCollection(iter(rows))
+
+    def __len__(self):
+        return len(self._all_rows)
+
+    def export(self, format, **kwargs):
+        """Export the RecordCollection to a given format (courtesy of Tablib)."""
+        return self.dataset.export(format, **kwargs)
+
+    @property
+    def dataset(self):
+        """A Tablib Dataset representation of the RecordCollection."""
+        # Create a new Tablib Dataset.
+        data = tablib.Dataset()
+
+        # If the RecordCollection is empty, just return the empty set
+        # Check number of rows by typecasting to list
+        if len(list(self)) == 0:
+            return data
+
+        # Set the column names as headers on Tablib Dataset.
+        first = self[0]
+
+        data.headers = first.keys()
+        for row in self.all():
+            row = _reduce_datetimes(row.values())
+            data.append(row)
+
+        return data
+
+    def all(self, as_dict=False, as_ordereddict=False):
+        """Returns a list of all rows for the RecordCollection. If they haven't
+        been fetched yet, consume the iterator and cache the results."""
+
+        # By calling list it calls the __iter__ method
+        rows = list(self)
+
+        if as_dict:
+            return [r.as_dict() for r in rows]
+        elif as_ordereddict:
+            return [r.as_dict(ordered=True) for r in rows]
+
+        return rows
+
+    def as_dict(self, ordered=False):
+        return self.all(as_dict=not (ordered), as_ordereddict=ordered)
+
+    def first(self, default=None, as_dict=False, as_ordereddict=False):
+        """Returns a single record for the RecordCollection, or `default`. If
+        `default` is an instance or subclass of Exception, then raise it
+        instead of returning it."""
+
+        # Try to get a record, or return/raise default.
+        try:
+            record = self[0]
+        except IndexError:
+            if is_exception(default):
+                raise default
+            return default
+
+        # Cast and return.
+        if as_dict:
+            return record.as_dict()
+        elif as_ordereddict:
+            return record.as_dict(ordered=True)
+        else:
+            return record
+
+    def one(self, default=None, as_dict=False, as_ordereddict=False):
+        """Returns a single record for the RecordCollection, ensuring that it
+        is the only record, or returns `default`. If `default` is an instance
+        or subclass of Exception, then raise it instead of returning it."""
+
+        # Ensure that we don't have more than one row.
+        try:
+            return self[1]
+        except IndexError:
+            return self.first(default=default, as_dict=as_dict,
+                              as_ordereddict=as_ordereddict)
+        else:
+            raise ValueError('RecordCollection contained more than one row. '
+                             'Expects only one row when using '
+                             'RecordCollection.one')
+
+    def scalar(self, default=None):
+        """Returns the first column of the first row, or `default`."""
+        row = self.one()
+        return row[0] if row else default
+
+
+class Database(object):
+    """A Database. Encapsulates a url and an SQLAlchemy engine with a pool of
+    connections.
+    """
+
+    def __init__(self, db_url=None, **kwargs):
+        # If no db_url was provided, fallback to $DATABASE_URL.
+        self.db_url = db_url or DATABASE_URL
+
+        if not self.db_url:
+            raise ValueError('You must provide a db_url.')
+
+        # Create an engine.
+        self._engine = create_engine(self.db_url, **kwargs)
+        self.open = True
+
+    def close(self):
+        """Closes the Database."""
+        self._engine.dispose()
+        self.open = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc, val, traceback):
+        self.close()
+
+    def __repr__(self):
+        return '<Database open={}>'.format(self.open)
+
+    def get_table_names(self):
+        """Returns a list of table names for the connected database."""
+
+        # Setup SQLAlchemy for Database inspection.
+        return inspect(self._engine).get_table_names()
+
+    def get_connection(self):
+        """Get a connection to this Database. Connections are retrieved from a
+        pool.
+        """
+        if not self.open:
+            raise exc.ResourceClosedError('Database closed.')
+
+        return Connection(self._engine.connect())
+
+    def query(self, query, fetchall=False, **params):
+        """Executes the given SQL query against the Database. Parameters can,
+        optionally, be provided. Returns a RecordCollection, which can be
+        iterated over to get result rows as dictionaries.
+        """
+        with self.get_connection() as conn:
+            return conn.query(query, fetchall, **params)
+
+    def bulk_query(self, query, *multiparams):
+        """Bulk insert or update."""
+
+        with self.get_connection() as conn:
+            conn.bulk_query(query, *multiparams)
+
+
+class Connection(object):
+    """A Database connection."""
+
+    def __init__(self, connection):
+        self._conn = connection
+        self.open = not connection.closed
+
+    def close(self):
+        self._conn.close()
+        self.open = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc, val, traceback):
+        self.close()
+
+    def __repr__(self):
+        return '<Connection open={}>'.format(self.open)
+
+    def query(self, query, fetchall=False, **params):
+        """Executes the given SQL query against the connected Database.
+        Parameters can, optionally, be provided. Returns a RecordCollection,
+        which can be iterated over to get result rows as dictionaries.
+        """
+
+        # Execute the given query.
+        cursor = self._conn.execute(text(query), **params)  # TODO: PARAMS GO HERE
+
+        # Row-by-row Record generator.
+        row_gen = (Record(cursor.keys(), row) for row in cursor)
+
+        # Convert psycopg2 results to RecordCollection.
+        results = RecordCollection(row_gen)
+
+        # Fetch all results if desired.
+        if fetchall:
+            results.all()
+
+        return results
+
+    def bulk_query(self, query, *multiparams):
+        """Bulk insert or update."""
+
+        self._conn.execute(text(query), *multiparams)
+
+
+def _reduce_datetimes(row):
+    """Receives a row, converts datetimes to strings."""
+
+    row = list(row)
+
+    for i in range(len(row)):
+        if hasattr(row[i], 'isoformat'):
+            row[i] = row[i].isoformat()
+    return tuple(row)

OneForAll/common/request.py

@@ -0,0 +1,264 @@
+import json
+from threading import Thread
+from queue import Queue
+
+import tqdm
+import requests
+from bs4 import BeautifulSoup
+
+from common import utils
+from config.log import logger
+from common.database import Database
+from config import settings
+
+
+def req_thread_count():
+    count = settings.request_thread_count
+    if isinstance(count, int):
+        count = max(16, count)
+    else:
+        count = utils.get_request_count()
+    logger.log('DEBUG', f'Number of request threads {count}')
+    return count
+
+
+def get_port_seq(port):
+    logger.log('DEBUG', 'Getting port range')
+    ports = set()
+    if isinstance(port, (set, list, tuple)):
+        ports = port
+    elif isinstance(port, int):
+        if 0 <= port <= 65535:
+            ports = {port}
+    elif port in {'small', 'medium', 'large'}:
+        logger.log('DEBUG', f'{port} port range')
+        ports = settings.ports.get(port)
+    if not ports:  # 意外情况
+        logger.log('ERROR', 'The specified request port range is incorrect')
+        ports = {80}
+    logger.log('INFOR', f'Port range:{ports}')
+    return set(ports)
+
+
+def gen_req_url(domain, port):
+    if str(port).endswith('443'):
+        url = f'https://{domain}:{port}'
+        if port == 443:
+            url = f'https://{domain}'
+        return url
+    url = f'http://{domain}:{port}'
+    if port == 80:
+        url = f'http://{domain}'
+    return url
+
+
+def gen_req_data(data, ports):
+    logger.log('INFOR', 'Generating request urls')
+    req_data = list()
+    req_urls = set()
+    for info in data:
+        resolve = info.get('resolve')
+        # 解析不成功的子域不进行http请求探测
+        if resolve != 1:
+            continue
+        subdomain = info.get('subdomain')
+        for port in ports:
+            tmp_info = info.copy()
+            tmp_info['port'] = port
+            url = gen_req_url(subdomain, port)
+            tmp_info['url'] = url
+            req_data.append(tmp_info)
+            req_urls.add(url)
+    return req_data, req_urls
+
+
+def get_html_title(markup):
+    """
+    获取标题
+
+    :param markup: html标签
+    :return: 标题
+    """
+    soup = BeautifulSoup(markup, 'html.parser')
+
+    title = soup.title
+    if title:
+        return title.text
+
+    h1 = soup.h1
+    if h1:
+        return h1.text
+
+    h2 = soup.h2
+    if h2:
+        return h2.text
+
+    h3 = soup.h3
+    if h3:
+        return h3.text
+
+    desc = soup.find('meta', attrs={'name': 'description'})
+    if desc:
+        return desc['content']
+
+    word = soup.find('meta', attrs={'name': 'keywords'})
+    if word:
+        return word['content']
+
+    text = soup.text
+    if len(text) <= 200:
+        return repr(text)
+
+    return 'None'
+
+
+def get_jump_urls(history):
+    urls = list()
+    for resp in history:
+        urls.append(str(resp.url))
+    return urls
+
+
+def get_progress_bar(total):
+    bar = tqdm.tqdm()
+    bar.total = total
+    bar.desc = 'Request Progress'
+    bar.ncols = 80
+    return bar
+
+
+def get_resp(url, session):
+    timeout = settings.request_timeout_second
+    redirect = settings.request_allow_redirect
+    proxy = utils.get_proxy()
+    try:
+        resp = session.get(url, timeout=timeout, allow_redirects=redirect, proxies=proxy)
+    except Exception as e:
+        logger.log('DEBUG', e.args)
+        resp = e
+    return resp
+
+
+def request(urls_queue, resp_queue, session):
+    while not urls_queue.empty():
+        index, url = urls_queue.get()
+        resp = get_resp(url, session)
+        resp_queue.put((index, resp))
+        urls_queue.task_done()
+
+
+def progress(bar, total, urls_queue):
+    while True:
+        remaining = urls_queue.qsize()
+        done = total - remaining
+        bar.n = done
+        bar.update()
+        if remaining == 0:
+            break
+
+
+def get_session():
+    header = utils.gen_fake_header()
+    verify = settings.request_ssl_verify
+    redirect_limit = settings.request_redirect_limit
+    session = requests.Session()
+    session.trust_env = False
+    session.headers = header
+    session.verify = verify
+    session.max_redirects = redirect_limit
+    return session
+
+
+def gen_new_info(info, resp):
+    if isinstance(resp, Exception):
+        info['reason'] = str(resp.args)
+        info['request'] = 0
+        info['alive'] = 0
+        return info
+    info['reason'] = resp.reason
+    code = resp.status_code
+    info['status'] = code
+    info['request'] = 1
+    if code == 400 or code >= 500:
+        info['alive'] = 0
+    else:
+        info['alive'] = 1
+    headers = resp.headers
+    if settings.enable_banner_identify:
+        info['banner'] = utils.get_sample_banner(headers)
+    info['header'] = json.dumps(dict(headers))
+    history = resp.history
+    info['history'] = json.dumps(get_jump_urls(history))
+    text = utils.decode_resp_text(resp)
+    title = get_html_title(text).strip()
+    info['title'] = utils.remove_invalid_string(title)
+    info['response'] = utils.remove_invalid_string(text)
+    return info
+
+
+def save(name, total, req_data, resp_queue):
+    db = Database()
+    db.create_table(name)
+    i = 0
+    while True:
+        if not resp_queue.empty():
+            i += 1
+            index, resp = resp_queue.get()
+            old_info = req_data[index]
+            new_info = gen_new_info(old_info, resp)
+            db.insert_table(name, new_info)
+            resp_queue.task_done()
+        if i >= total:  # 得存入完所有请求结果才能结束
+            break
+    db.close()
+
+
+def bulk_request(domain, req_data, ret=False):
+    logger.log('INFOR', 'Requesting urls in bulk')
+    resp_queue = Queue()
+    urls_queue = Queue()
+    task_count = len(req_data)
+    for index, info in enumerate(req_data):
+        url = info.get('url')
+        urls_queue.put((index, url))
+    session = get_session()
+    thread_count = req_thread_count()
+    if task_count <= thread_count:
+        # 如果请求任务数很小不用创建很多线程了
+        thread_count = task_count
+    bar = get_progress_bar(task_count)
+
+    progress_thread = Thread(target=progress, name='ProgressThread',
+                             args=(bar, task_count, urls_queue), daemon=True)
+    progress_thread.start()
+
+    for i in range(thread_count):
+        request_thread = Thread(target=request, name=f'RequestThread-{i}',
+                                args=(urls_queue, resp_queue, session), daemon=True)
+        request_thread.start()
+    if ret:
+        urls_queue.join()
+        return resp_queue
+    save_thread = Thread(target=save, name=f'SaveThread',
+                         args=(domain, task_count, req_data, resp_queue), daemon=True)
+    save_thread.start()
+    urls_queue.join()
+    save_thread.join()
+
+
+def run_request(domain, data, port):
+    """
+    HTTP request entrance
+
+    :param  str domain: domain to be requested
+    :param  list data: subdomains data to be requested
+    :param  any port: range of ports to be requested
+    :return list: result
+    """
+    logger.log('INFOR', f'Start requesting subdomains of {domain}')
+    data = utils.set_id_none(data)
+    ports = get_port_seq(port)
+    req_data, req_urls = gen_req_data(data, ports)
+    bulk_request(domain, req_data)
+    count = utils.count_alive(domain)
+    logger.log('INFOR', f'Found that {domain} has {count} alive subdomains')