mirror of https://github.com/coddrago/Heroku
commit
bde62d52d9
|
@ -1,11 +0,0 @@
|
||||||
import typing
|
|
||||||
|
|
||||||
class BaseTunnel:
|
|
||||||
async def start(self):
|
|
||||||
raise NotImplementedError("Subclasses must implement the 'start' method.")
|
|
||||||
|
|
||||||
async def stop(self):
|
|
||||||
raise NotImplementedError("Subclasses must implement the 'stop' method.")
|
|
||||||
|
|
||||||
async def wait_for_url(self, timeout: float) -> typing.Optional[str]:
|
|
||||||
raise NotImplementedError("Subclasses must implement the 'wait_for_url' method.")
|
|
|
@ -1,62 +0,0 @@
|
||||||
import typing
|
|
||||||
import logging
|
|
||||||
import asyncio
|
|
||||||
import contextvars
|
|
||||||
import functools
|
|
||||||
|
|
||||||
from pycloudflared import try_cloudflare
|
|
||||||
|
|
||||||
from .base_tunnel import BaseTunnel
|
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
|
||||||
|
|
||||||
|
|
||||||
class CloudflareTunnel(BaseTunnel):
|
|
||||||
def __init__(
|
|
||||||
self,
|
|
||||||
port: int,
|
|
||||||
verbose: bool = False,
|
|
||||||
change_url_callback: typing.Callable[[str], None] = None,
|
|
||||||
):
|
|
||||||
self.port = port
|
|
||||||
self.verbose = verbose
|
|
||||||
self._change_url_callback = change_url_callback
|
|
||||||
self._tunnel_url = None
|
|
||||||
self._url_available = asyncio.Event()
|
|
||||||
self._url_available.clear()
|
|
||||||
|
|
||||||
# to support python 3.8...
|
|
||||||
async def to_thread(self, func, /, *args, **kwargs):
|
|
||||||
loop = asyncio.get_running_loop()
|
|
||||||
ctx = contextvars.copy_context()
|
|
||||||
func_call = functools.partial(ctx.run, func, *args, **kwargs)
|
|
||||||
return await loop.run_in_executor(None, func_call)
|
|
||||||
|
|
||||||
async def start(self):
|
|
||||||
logger.debug(f"Attempting Cloudflare tunnel on port {self.port}...")
|
|
||||||
|
|
||||||
try:
|
|
||||||
self._tunnel_url = (await self.to_thread(try_cloudflare, port=self.port, verbose=self.verbose)).tunnel
|
|
||||||
logger.debug(f"Cloudflare tunnel established: {self._tunnel_url}")
|
|
||||||
|
|
||||||
if self._change_url_callback:
|
|
||||||
self._change_url_callback(self._tunnel_url)
|
|
||||||
|
|
||||||
self._url_available.set()
|
|
||||||
|
|
||||||
except Exception as e:
|
|
||||||
logger.error(f"Failed to establish Cloudflare tunnel: {e}")
|
|
||||||
raise
|
|
||||||
|
|
||||||
async def stop(self):
|
|
||||||
logger.debug("Stopping Cloudflare tunnel...")
|
|
||||||
try_cloudflare.terminate(self.port)
|
|
||||||
|
|
||||||
async def wait_for_url(self, timeout: float) -> typing.Optional[str]:
|
|
||||||
try:
|
|
||||||
await asyncio.wait_for(self._url_available.wait(), timeout)
|
|
||||||
return self._tunnel_url
|
|
||||||
except asyncio.TimeoutError:
|
|
||||||
logger.warning("Timeout waiting for Cloudflare URL.")
|
|
||||||
return None
|
|
|
@ -8,7 +8,6 @@ import os
|
||||||
import logging
|
import logging
|
||||||
import typing
|
import typing
|
||||||
from .ssh_tunnel import SSHTunnel
|
from .ssh_tunnel import SSHTunnel
|
||||||
from .cloudflare_tunnel import CloudflareTunnel
|
|
||||||
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
@ -27,7 +26,6 @@ class ProxyPasser:
|
||||||
self._verbose = verbose
|
self._verbose = verbose
|
||||||
self._tunnels = [
|
self._tunnels = [
|
||||||
SSHTunnel(port=port, change_url_callback=self._on_url_change),
|
SSHTunnel(port=port, change_url_callback=self._on_url_change),
|
||||||
CloudflareTunnel(port=port, verbose=verbose, change_url_callback=self._on_url_change)
|
|
||||||
]
|
]
|
||||||
|
|
||||||
|
|
||||||
|
|
|
@ -3,11 +3,10 @@ import logging
|
||||||
import asyncio
|
import asyncio
|
||||||
import re
|
import re
|
||||||
|
|
||||||
from .base_tunnel import BaseTunnel
|
|
||||||
|
|
||||||
logger = logging.getLogger(__name__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
class SSHTunnel(BaseTunnel):
|
class SSHTunnel():
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
port: int,
|
port: int,
|
||||||
|
|
Loading…
Reference in New Issue