Fork project

utils/routers/__init__.py

@@ -0,0 +1,34 @@
+from CTFd.utils import get_config
+
+from .frp import FrpRouter
+from .trp import TrpRouter
+
+_routers = {
+    'frp': FrpRouter,
+    'trp': TrpRouter,
+}
+
+
+def instanciate(cls):
+    return cls()
+
+
+@instanciate
+class Router:
+    _name = ''
+    _router = None
+
+    def __getattr__(self, name: str):
+        router_conftype = get_config("whale:router_type", "frp")
+        if Router._name != router_conftype:
+            Router._router = _routers[router_conftype]()
+            Router._name = router_conftype
+        return getattr(Router._router, name)
+
+    @staticmethod
+    def reset():
+        Router._name = ''
+        Router._router = None
+
+
+__all__ = ["Router"]

utils/routers/base.py

@@ -0,0 +1,25 @@
+import typing
+
+from ...models import WhaleContainer
+
+
+class BaseRouter:
+    name = None
+
+    def __init__(self):
+        pass
+
+    def access(self, container: WhaleContainer):
+        pass
+
+    def register(self, container: WhaleContainer):
+        pass
+
+    def unregister(self, container: WhaleContainer):
+        pass
+
+    def reload(self):
+        pass
+
+    def check_availability(self) -> typing.Tuple[bool, str]:
+        pass

utils/routers/frp.py

@@ -0,0 +1,132 @@
+import warnings
+
+from flask import current_app
+from requests import session, RequestException
+
+from CTFd.models import db
+from CTFd.utils import get_config, set_config, logging
+
+from .base import BaseRouter
+from ..cache import CacheProvider
+from ..db import DBContainer
+from ..exceptions import WhaleError, WhaleWarning
+from ...models import WhaleContainer
+
+
+class FrpRouter(BaseRouter):
+    name = "frp"
+    types = {
+        'direct': 'tcp',
+        'http': 'http',
+    }
+
+    class FrpRule:
+        def __init__(self, name, config):
+            self.name = name
+            self.config = config
+
+        def __str__(self) -> str:
+            return f'[{self.name}]\n' + '\n'.join(f'{k} = {v}' for k, v in self.config.items())
+
+    def __init__(self):
+        super().__init__()
+        self.ses = session()
+        self.url = get_config("whale:frp_api_url").rstrip("/")
+        self.common = ''
+        try:
+            CacheProvider(app=current_app).init_port_sets()
+        except Exception:
+            warnings.warn(
+                "cache initialization failed",
+                WhaleWarning
+            )
+
+    def reload(self, exclude=None):
+        rules = []
+        for container in DBContainer.get_all_alive_container():
+            if container.uuid == exclude:
+                continue
+            name = f'{container.challenge.redirect_type}_{container.user_id}_{container.uuid}'
+            config = {
+                'type': self.types[container.challenge.redirect_type],
+                'local_ip': f'{container.user_id}-{container.uuid}',
+                'local_port': container.challenge.redirect_port,
+                'use_compression': 'true',
+            }
+            if config['type'] == 'http':
+                config['subdomain'] = container.http_subdomain
+            elif config['type'] == 'tcp':
+                config['remote_port'] = container.port
+            rules.append(self.FrpRule(name, config))
+
+        try:
+            if not self.common:
+                common = get_config("whale:frp_config_template", '')
+                if '[common]' in common:
+                    self.common = common
+                else:
+                    remote = self.ses.get(f'{self.url}/api/config')
+                    assert remote.status_code == 200
+                    set_config("whale:frp_config_template", remote.text)
+                    self.common = remote.text
+            config = self.common + '\n' + '\n'.join(str(r) for r in rules)
+            assert self.ses.put(
+                f'{self.url}/api/config', config, timeout=5
+            ).status_code == 200
+            assert self.ses.get(
+                f'{self.url}/api/reload', timeout=5
+            ).status_code == 200
+        except (RequestException, AssertionError) as e:
+            raise WhaleError(
+                '\nfrpc request failed\n' +
+                (f'{e}\n' if str(e) else '') +
+                'please check the frp related configs'
+            ) from None
+
+    def access(self, container: WhaleContainer):
+        if container.challenge.redirect_type == 'direct':
+            return f'nc {get_config("whale:frp_direct_ip_address", "127.0.0.1")} {container.port}'
+        elif container.challenge.redirect_type == 'http':
+            host = get_config("whale:frp_http_domain_suffix", "")
+            port = get_config("whale:frp_http_port", "80")
+            host += f':{port}' if port != 80 else ''
+            return f'<a target="_blank" href="http://{container.http_subdomain}.{host}/">Link to the Challenge</a>'
+        return ''
+
+    def register(self, container: WhaleContainer):
+        if container.challenge.redirect_type == 'direct':
+            if not container.port:
+                port = CacheProvider(app=current_app).get_available_port()
+                if not port:
+                    return False, 'No available ports. Please wait for a few minutes.'
+                container.port = port
+                db.session.commit()
+        elif container.challenge.redirect_type == 'http':
+            # config['subdomain'] = container.http_subdomain
+            pass
+        self.reload()
+        return True, 'success'
+
+    def unregister(self, container: WhaleContainer):
+        if container.challenge.redirect_type == 'direct':
+            try:
+                redis_util = CacheProvider(app=current_app)
+                redis_util.add_available_port(container.port)
+            except Exception as e:
+                logging.log(
+                    'whale', 'Error deleting port from cache',
+                    name=container.user.name,
+                    challenge_id=container.challenge_id,
+                )
+                return False, 'Error deleting port from cache'
+        self.reload(exclude=container.uuid)
+        return True, 'success'
+
+    def check_availability(self):
+        try:
+            resp = self.ses.get(f'{self.url}/api/status', timeout=2.0)
+        except RequestException as e:
+            return False, 'Unable to access frpc admin api'
+        if resp.status_code == 401:
+            return False, 'frpc admin api unauthorized'
+        return True, 'Available'

utils/routers/trp.py

@@ -0,0 +1,69 @@
+import traceback
+from requests import session, RequestException, HTTPError
+
+from CTFd.utils import get_config
+from .base import BaseRouter
+from ..db import DBContainer, WhaleContainer
+
+
+class TrpRouter(BaseRouter):
+    name = "trp"
+
+    def __init__(self):
+        super().__init__()
+        self.ses = session()
+        self.url = get_config('whale:trp_api_url', '').rstrip("/")
+        self.common = ''
+        for container in DBContainer.get_all_alive_container():
+            self.register(container)
+
+    @staticmethod
+    def get_domain(container: WhaleContainer):
+        domain = get_config('whale:trp_domain_suffix', '127.0.0.1.nip.io').lstrip('.')
+        domain = f'{container.uuid}.{domain}'
+        return domain
+
+    def access(self, container: WhaleContainer):
+        ch_type = container.challenge.redirect_type
+        domain = self.get_domain(container)
+        port = get_config('whale:trp_listening_port', 1443)
+        if ch_type == 'direct':
+            return f'from pwn import *<br>remote("{domain}", {port}, ssl=True).interactive()'
+        elif ch_type == 'http':
+            return f'https://{domain}' + (f':{port}' if port != 443 else '')
+        else:
+            return f'[ssl] {domain} {port}'
+
+    def register(self, container: WhaleContainer):
+        try:
+            resp = self.ses.post(f'{self.url}/rule/{self.get_domain(container)}', json={
+                'target': f'{container.user_id}-{container.uuid}:{container.challenge.redirect_port}',
+                'source': None,
+            })
+            resp.raise_for_status()
+            return True, 'success'
+        except HTTPError as e:
+            return False, e.response.text
+        except RequestException as e:
+            print(traceback.format_exc())
+            return False, 'unable to access trp Api'
+
+    def unregister(self, container: WhaleContainer):
+        try:
+            resp = self.ses.delete(f'{self.url}/rule/{self.get_domain(container)}')
+            resp.raise_for_status()
+            return True, 'success'
+        except HTTPError as e:
+            return False, e.response.text
+        except RequestException as e:
+            print(traceback.format_exc())
+            return False, 'unable to access trp Api'
+
+    def check_availability(self):
+        try:
+            resp = self.ses.get(f'{self.url}/rules').json()
+        except RequestException as e:
+            return False, 'Unable to access trp admin api'
+        except Exception as e:
+            return False, 'Unknown trp error'
+        return True, 'Available'