mirror of https://github.com/coddrago/Heroku
101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
import asyncio
|
|
import logging
|
|
import re
|
|
import typing
|
|
|
|
from .. import utils
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ProxyPasser:
|
|
def __init__(self):
|
|
self._tunnel_url = None
|
|
self._sproc = None
|
|
self._url_available = asyncio.Event()
|
|
self._url_available.set()
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def _sleep_for_task(self, callback: callable, data: bytes, delay: int):
|
|
await asyncio.sleep(delay)
|
|
await callback(data.decode("utf-8"))
|
|
|
|
async def _read_stream(
|
|
self,
|
|
callback: callable,
|
|
stream: typing.BinaryIO,
|
|
delay: int,
|
|
) -> None:
|
|
last_task = None
|
|
for getline in iter(stream.readline, ""):
|
|
data_chunk = await getline
|
|
if not data_chunk:
|
|
if last_task:
|
|
last_task.cancel()
|
|
await callback(data_chunk.decode("utf-8"))
|
|
if not self._url_available.is_set():
|
|
self._url_available.set()
|
|
|
|
if last_task:
|
|
last_task.cancel()
|
|
|
|
last_task = asyncio.ensure_future(
|
|
self._sleep_for_task(callback, data_chunk, delay)
|
|
)
|
|
|
|
def kill(self):
|
|
try:
|
|
self._sproc.terminate()
|
|
except Exception:
|
|
logger.exception("Failed to kill proxy pass process")
|
|
else:
|
|
logger.debug("Proxy pass tunnel killed")
|
|
|
|
async def _process_stream(self, stdout_line: str) -> None:
|
|
regex = r"tunneled.*?(https:\/\/.*?\..*?\.[a-z]+)"
|
|
|
|
if re.search(regex, stdout_line):
|
|
logger.debug("Proxy pass tunneled: %s", stdout_line)
|
|
self._tunnel_url = re.search(regex, stdout_line)[1]
|
|
self._url_available.set()
|
|
|
|
async def get_url(self, port: int) -> typing.Optional[str]:
|
|
async with self._lock:
|
|
if self._tunnel_url:
|
|
try:
|
|
await asyncio.wait_for(self._sproc.wait(), timeout=0.1)
|
|
except asyncio.TimeoutError:
|
|
return self._tunnel_url
|
|
else:
|
|
self.kill()
|
|
|
|
logger.debug("Starting proxy pass shell")
|
|
self._sproc = await asyncio.create_subprocess_shell(
|
|
"ssh -o StrictHostKeyChecking=no -R"
|
|
f" 80:127.0.0.1:{port} nokey@localhost.run",
|
|
stdin=asyncio.subprocess.PIPE,
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
|
|
utils.atexit(self.kill)
|
|
|
|
self._url_available = asyncio.Event()
|
|
logger.debug("Starting proxy pass reader")
|
|
asyncio.ensure_future(
|
|
self._read_stream(
|
|
self._process_stream,
|
|
self._sproc.stdout,
|
|
1,
|
|
)
|
|
)
|
|
|
|
try:
|
|
await asyncio.wait_for(self._url_available.wait(), 15)
|
|
except asyncio.TimeoutError:
|
|
self.kill()
|
|
self._tunnel_url = None
|
|
return await self.get_url(port)
|
|
|
|
return self._tunnel_url
|