229 lines
7.4 KiB
Python
229 lines
7.4 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""
|
|
proxy.py
|
|
~~~~~~~~
|
|
⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on
|
|
Network monitoring, controls & Application development, testing, debugging.
|
|
|
|
:copyright: (c) 2013-present by Abhinav Singh and contributors.
|
|
:license: BSD, see LICENSE for more details.
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import logging
|
|
|
|
from typing import List, Optional, Any
|
|
|
|
from .core.acceptor import AcceptorPool, ThreadlessPool, Listener
|
|
from .core.event import EventManager
|
|
from .common.utils import bytes_
|
|
from .common.flag import FlagParser, flags
|
|
from .common.constants import DEFAULT_LOG_FILE, DEFAULT_LOG_FORMAT, DEFAULT_LOG_LEVEL
|
|
from .common.constants import DEFAULT_OPEN_FILE_LIMIT, DEFAULT_PLUGINS, DEFAULT_VERSION
|
|
from .common.constants import DEFAULT_ENABLE_DASHBOARD, DEFAULT_WORK_KLASS, DEFAULT_PID_FILE
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
flags.add_argument(
|
|
'--version',
|
|
'-v',
|
|
action='store_true',
|
|
default=DEFAULT_VERSION,
|
|
help='Prints proxy.py version.',
|
|
)
|
|
|
|
# TODO: Convert me into 1-letter choices
|
|
# TODO: Add --verbose option which also
|
|
# starts to log traffic flowing between
|
|
# clients and upstream servers.
|
|
flags.add_argument(
|
|
'--log-level',
|
|
type=str,
|
|
default=DEFAULT_LOG_LEVEL,
|
|
help='Valid options: DEBUG, INFO (default), WARNING, ERROR, CRITICAL. '
|
|
'Both upper and lowercase values are allowed. '
|
|
'You may also simply use the leading character e.g. --log-level d',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--log-file',
|
|
type=str,
|
|
default=DEFAULT_LOG_FILE,
|
|
help='Default: sys.stdout. Log file destination.',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--log-format',
|
|
type=str,
|
|
default=DEFAULT_LOG_FORMAT,
|
|
help='Log format for Python logger.',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--open-file-limit',
|
|
type=int,
|
|
default=DEFAULT_OPEN_FILE_LIMIT,
|
|
help='Default: 1024. Maximum number of files (TCP connections) '
|
|
'that proxy.py can open concurrently.',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--plugins',
|
|
action='append',
|
|
nargs='+',
|
|
default=DEFAULT_PLUGINS,
|
|
help='Comma separated plugins. ' +
|
|
'You may use --plugins flag multiple times.',
|
|
)
|
|
|
|
# TODO: Ideally all `--enable-*` flags must be at the top-level.
|
|
# --enable-dashboard is specially needed here because
|
|
# ProxyDashboard class is not imported anywhere.
|
|
#
|
|
# Due to which, if we move this flag definition within dashboard
|
|
# plugin, users will have to explicitly enable dashboard plugin
|
|
# to also use flags provided by it.
|
|
flags.add_argument(
|
|
'--enable-dashboard',
|
|
action='store_true',
|
|
default=DEFAULT_ENABLE_DASHBOARD,
|
|
help='Default: False. Enables proxy.py dashboard.',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--work-klass',
|
|
type=str,
|
|
default=DEFAULT_WORK_KLASS,
|
|
help='Default: ' + DEFAULT_WORK_KLASS +
|
|
'. Work klass to use for work execution.',
|
|
)
|
|
|
|
flags.add_argument(
|
|
'--pid-file',
|
|
type=str,
|
|
default=DEFAULT_PID_FILE,
|
|
help='Default: None. Save "parent" process ID to a file.',
|
|
)
|
|
|
|
|
|
class Proxy:
|
|
"""Proxy is a context manager to control proxy.py library core.
|
|
|
|
By default, :class:`~proxy.core.pool.AcceptorPool` is started with
|
|
:class:`~proxy.http.handler.HttpProtocolHandler` work class.
|
|
By definition, it expects HTTP traffic to flow between clients and server.
|
|
|
|
In ``--threadless`` mode and without ``--local-executor``,
|
|
a :class:`~proxy.core.executors.ThreadlessPool` is also started.
|
|
Executor pool receives newly accepted work by :class:`~proxy.core.acceptor.Acceptor`
|
|
and creates an instance of work class for processing the received work.
|
|
|
|
Optionally, Proxy class also initializes the EventManager.
|
|
A multi-process safe pubsub system which can be used to build various
|
|
patterns for message sharing and/or signaling.
|
|
"""
|
|
|
|
def __init__(self, input_args: Optional[List[str]] = None, **opts: Any) -> None:
|
|
self.flags = FlagParser.initialize(input_args, **opts)
|
|
self.listener: Optional[Listener] = None
|
|
self.executors: Optional[ThreadlessPool] = None
|
|
self.acceptors: Optional[AcceptorPool] = None
|
|
self.event_manager: Optional[EventManager] = None
|
|
|
|
def __enter__(self) -> 'Proxy':
|
|
self.setup()
|
|
return self
|
|
|
|
def __exit__(self, *args: Any) -> None:
|
|
self.shutdown()
|
|
|
|
def setup(self) -> None:
|
|
# TODO: Introduce cron feature
|
|
# https://github.com/abhinavsingh/proxy.py/issues/392
|
|
#
|
|
# TODO: Introduce ability to publish
|
|
# adhoc events which can modify behaviour of server
|
|
# at runtime. Example, updating flags, plugin
|
|
# configuration etc.
|
|
#
|
|
# TODO: Python shell within running proxy.py environment?
|
|
#
|
|
# TODO: Pid watcher which watches for processes started
|
|
# by proxy.py core. May be alert or restart those processes
|
|
# on failure.
|
|
self._write_pid_file()
|
|
# We setup listeners first because of flags.port override
|
|
# in case of ephemeral port being used
|
|
self.listener = Listener(flags=self.flags)
|
|
self.listener.setup()
|
|
# Override flags.port to match the actual port
|
|
# we are listening upon. This is necessary to preserve
|
|
# the server port when `--port=0` is used.
|
|
self.flags.port = self.listener._port
|
|
# Setup EventManager
|
|
if self.flags.enable_events:
|
|
logger.info('Core Event enabled')
|
|
self.event_manager = EventManager()
|
|
self.event_manager.setup()
|
|
event_queue = self.event_manager.queue \
|
|
if self.event_manager is not None \
|
|
else None
|
|
# Setup remote executors
|
|
if not self.flags.local_executor:
|
|
self.executors = ThreadlessPool(
|
|
flags=self.flags,
|
|
event_queue=event_queue,
|
|
)
|
|
self.executors.setup()
|
|
# Setup acceptors
|
|
self.acceptors = AcceptorPool(
|
|
flags=self.flags,
|
|
listener=self.listener,
|
|
executor_queues=self.executors.work_queues if self.executors else [],
|
|
executor_pids=self.executors.work_pids if self.executors else [],
|
|
executor_locks=self.executors.work_locks if self.executors else [],
|
|
event_queue=event_queue,
|
|
)
|
|
self.acceptors.setup()
|
|
# TODO: May be close listener fd as we don't need it now
|
|
|
|
def shutdown(self) -> None:
|
|
assert self.acceptors
|
|
self.acceptors.shutdown()
|
|
if not self.flags.local_executor:
|
|
assert self.executors
|
|
self.executors.shutdown()
|
|
if self.flags.enable_events:
|
|
assert self.event_manager is not None
|
|
self.event_manager.shutdown()
|
|
assert self.listener
|
|
self.listener.shutdown()
|
|
self._delete_pid_file()
|
|
|
|
def _write_pid_file(self) -> None:
|
|
if self.flags.pid_file is not None:
|
|
# NOTE: Multiple instances of proxy.py running on
|
|
# same host machine will currently result in overwriting the PID file
|
|
with open(self.flags.pid_file, 'wb') as pid_file:
|
|
pid_file.write(bytes_(os.getpid()))
|
|
|
|
def _delete_pid_file(self) -> None:
|
|
if self.flags.pid_file and os.path.exists(self.flags.pid_file):
|
|
os.remove(self.flags.pid_file)
|
|
|
|
|
|
def main(**opts: Any) -> None:
|
|
try:
|
|
with Proxy(sys.argv[1:], **opts):
|
|
while True:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
pass
|
|
|
|
|
|
def entry_point() -> None:
|
|
main()
|