From 3861310d71d28b7aefa438b12698c7b7f0649890 Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Sun, 13 Jun 2021 17:15:22 -0400
Subject: [PATCH] Even better file close handling

For linux, we used to send one or two EOFs via C-d to the process to
signal exit, however this was inconsistent. Depending on the previous
input from the attacker, sometimes one was needed, sometimes two.
Sometimes, we even observed more than two being needed. Instead, we now
simply loop sending one at a time and checking for the end delimeter.
This will be slightly slower, but avoids hangups or accidentally closing
the shell.
---
 pwncat/platform/linux.py | 28 +++++++++++++++-------------
 tests/test_platform.py   |  1 +
 2 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/pwncat/platform/linux.py b/pwncat/platform/linux.py
index 54c09de..21f69a7 100644
--- a/pwncat/platform/linux.py
+++ b/pwncat/platform/linux.py
@@ -458,19 +458,21 @@ class LinuxWriter(BufferedIOBase):
             self.detach()
             return
 
-        # Indicate EOF
-        self.popen.stdin.write(b"\x04")
-        if self.since_newline != 0:
-            self.popen.stdin.write(b"\x04")
-
-        try:
-            # Check for completion
-            self.popen.wait(timeout=100)
-        except pwncat.subprocess.TimeoutExpired:
-            # Nope, force terminate with C-c
-            # self.popen.terminate()
-            # Cleanup
-            self.popen.wait()
+        # The number of C-d's needed to trigger an EOF in
+        # the process and exit is inconsistent based on the
+        # previous input. So, instead of trying to be deterministic,
+        # we simply send one and check. We do this until we find
+        # the ending delimeter and then exit. If the `on_close`
+        # hook was setup properly, this should be fine.
+        while True:
+            try:
+                self.popen.stdin.write(b"\x04")
+                self.popen.stdin.flush()
+                # Check for completion
+                self.popen.wait(timeout=0.1)
+                break
+            except pwncat.subprocess.TimeoutExpired:
+                continue
 
         # Ensure we don't touch stdio again
         self.detach()
diff --git a/tests/test_platform.py b/tests/test_platform.py
index efa5ae5..cc1ddcc 100644
--- a/tests/test_platform.py
+++ b/tests/test_platform.py
@@ -5,6 +5,7 @@ import base64
 import subprocess
 
 import pytest
+
 from pwncat.util import random_string
 from pwncat.platform.windows import PowershellError