Saverr/pwncat
1
0
Fork 0
mirror of https://github.com/calebstewart/pwncat.git synced 2024-11-30 20:34:15 +01:00
Code Projects Releases Wiki Activity

Added exception handling for state transition

Browse Source
This commit is contained in:
Caleb Stewart 2021-06-16 18:48:10 -04:00
parent 2691f77c76
commit 07be104ddd
1 changed files with 32 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
pwncat/manager.py
View File

@ -36,9 +36,9 @@ import pwncat.modules.enumerate
from pwncat.util import RawModeExit, console from pwncat.util import RawModeExit, console
from pwncat.config import Config from pwncat.config import Config
from pwncat.target import Target from pwncat.target import Target
from pwncat.channel import Channel, ChannelClosed from pwncat.channel import Channel, ChannelError, ChannelClosed
from pwncat.commands import CommandParser from pwncat.commands import CommandParser
from pwncat.platform import Platform from pwncat.platform import Platform, PlatformError
class InteractiveExit(Exception): class InteractiveExit(Exception):
@ -319,9 +319,13 @@ class Session:
while self.layers: while self.layers:
self.layers.pop()(self) self.layers.pop()(self)
try:
self.platform.exit() self.platform.exit()
self.platform.channel.close() self.platform.channel.close()
except (PlatformError, ChannelError) as exc:
self.log(
f"[yellow]warning[/yellow]: unexpected exception while closing: {exc}"
)
self.died() self.died()
@ -574,15 +578,15 @@ class Manager:
interactive_complete = threading.Event() interactive_complete = threading.Event()
def output_thread_main(): def output_thread_main(target: Session):
while not interactive_complete.is_set(): while not interactive_complete.is_set():
try:
data = self.target.platform.channel.recv(4096) data = target.platform.channel.recv(4096)
if data != b"" and data is not None: if data != b"" and data is not None:
try: try:
data = self.target.platform.process_output(data) data = target.platform.process_output(data)
sys.stdout.buffer.write(data) sys.stdout.buffer.write(data)
sys.stdout.buffer.flush() sys.stdout.buffer.flush()
except RawModeExit: except RawModeExit:
@ -590,20 +594,25 @@ class Manager:
else: else:
interactive_complete.wait(timeout=0.1) interactive_complete.wait(timeout=0.1)
output_thread = threading.Thread(target=output_thread_main) except ChannelError:
interactive_complete.set()
output_thread = threading.Thread(
target=output_thread_main, args=[self.target]
)
output_thread.start() output_thread.start()
channel_closed = False try:
try: try:
self.target.platform.interactive_loop(interactive_complete) self.target.platform.interactive_loop(interactive_complete)
except RawModeExit: except RawModeExit:
pass pass
self.target.platform.interactive = False
except ChannelClosed: except ChannelClosed:
channel_closed = True
self.log( self.log(
f"[yellow]warning[/yellow]: {self.target.platform}: connection reset" f"[yellow]warning[/yellow]: {self.target.platform}: connection reset"
) )
self.target.died()
except Exception: except Exception:
pwncat.util.console.print_exception() pwncat.util.console.print_exception()
@ -611,12 +620,6 @@ class Manager:
interactive_complete.set() interactive_complete.set()
output_thread.join() output_thread.join()
# Exit interactive mode
if channel_closed:
self.target.died()
else:
self.target.platform.interactive = False
def create_session(self, platform: str, channel: Channel = None, **kwargs): def create_session(self, platform: str, channel: Channel = None, **kwargs):
""" """
Open a new session from a new or existing platform. If the platform Open a new session from a new or existing platform. If the platform