From 5544954852945f55091964afdb2f0d511197f413 Mon Sep 17 00:00:00 2001
From: Caleb Stewart <caleb.stewart94@gmail.com>
Date: Fri, 18 Jun 2021 14:10:56 -0400
Subject: [PATCH] Added better file io tests which pass

---
 tests/test_fileio.py   | 49 ++++++++++++++++++++++++++++++++++++++++++
 tests/test_platform.py | 27 ++++-------------------
 2 files changed, 53 insertions(+), 23 deletions(-)
 create mode 100644 tests/test_fileio.py

diff --git a/tests/test_fileio.py b/tests/test_fileio.py
new file mode 100644
index 0000000..e1b40a8
--- /dev/null
+++ b/tests/test_fileio.py
@@ -0,0 +1,49 @@
+#!/usr/bin/env python3
+
+from pwncat.util import random_string
+
+
+def do_file_test(session, content):
+    """Do a generic file test"""
+
+    name = random_string() + ".txt"
+    mode = "b" if isinstance(content, bytes) else ""
+
+    with session.platform.open(name, mode + "w") as filp:
+        assert filp.write(content) == len(content)
+
+    with session.platform.open(name, mode + "r") as filp:
+        assert filp.read() == content
+
+    # In some cases, the act of reading/writing causes a shell to hang
+    # so double check that.
+    result = session.platform.run(
+        ["echo", "hello world"], capture_output=True, text=True
+    )
+    assert result.stdout == "hello world\n"
+
+
+def test_small_text(session):
+    """Test writing a small text-only file"""
+
+    do_file_test(session, "hello world")
+
+
+def test_large_text(session):
+    """Test writing and reading a large text file"""
+
+    contents = ("A" * 1000 + "\n") * 10
+    do_file_test(session, contents)
+
+
+def test_small_binary(session):
+    """Test writing a small amount of binary data"""
+
+    contents = bytes(list(range(32)))
+    do_file_test(session, contents)
+
+
+def test_large_binary(session):
+
+    contents = bytes(list(range(32))) * 400
+    do_file_test(session, contents)
diff --git a/tests/test_platform.py b/tests/test_platform.py
index cc1ddcc..c81525e 100644
--- a/tests/test_platform.py
+++ b/tests/test_platform.py
@@ -10,27 +10,8 @@ from pwncat.util import random_string
 from pwncat.platform.windows import PowershellError
 
 
-def test_platform_file_io(session):
-    """ Test file read/write of printable data """
-
-    # Generate random binary data
-    contents = os.urandom(1024)
-
-    # Create a new temporary file
-    with session.platform.tempfile(mode="wb") as filp:
-        filp.write(contents)
-        path = filp.name
-
-    # Ensure it exists
-    assert session.platform.Path(path).exists()
-
-    # Read the data back and ensure it matches
-    with session.platform.open(path, "rb") as filp:
-        assert contents == filp.read()
-
-
 def test_platform_dir_io(session):
-    """ Test creating a directory and interacting with the contents """
+    """Test creating a directory and interacting with the contents"""
 
     # Create a path object representing the new remote directory
     path = session.platform.Path(random_string())
@@ -61,7 +42,7 @@ def test_platform_run(session):
 
 
 def test_platform_su(session):
-    """ Test running `su` """
+    """Test running `su`"""
 
     try:
         session.platform.su("john", "P@ssw0rd")
@@ -77,7 +58,7 @@ def test_platform_su(session):
 
 
 def test_platform_sudo(session):
-    """ Testing running `sudo` """
+    """Testing running `sudo`"""
 
     try:
 
@@ -103,7 +84,7 @@ def test_platform_sudo(session):
 
 
 def test_windows_powershell(windows):
-    """ Test powershell execution """
+    """Test powershell execution"""
 
     # Run a real powershell snippet
     r = windows.platform.powershell("$PSVersionTable.PSVersion")