From 5b0db6b7739b683ade721fa3f1ec555b1769eb03 Mon Sep 17 00:00:00 2001 From: Tushar Sadhwani Date: Tue, 21 Feb 2023 13:23:23 +0530 Subject: [PATCH] Make CLI arguments accessible in zxpy programs (#52) * Fix sys.argv passed to tests * Improve test assertion * Add $ arg support * Normalize /bin/sh as first argument of script_args * Add type def --- tests/test_files/argv.py | 8 +++- tests/test_files/injection.py | 6 +++ tests/zxpy_test.py | 60 ++++++++++++++++++++++++++- zx.py | 78 ++++++++++++++++++++++++++++++++++- 4 files changed, 147 insertions(+), 5 deletions(-) create mode 100644 tests/test_files/injection.py diff --git a/tests/test_files/argv.py b/tests/test_files/argv.py index 27a09ea..5f0916f 100644 --- a/tests/test_files/argv.py +++ b/tests/test_files/argv.py @@ -1,4 +1,8 @@ import sys -assert len(sys.argv) == 1 -assert sys.argv[0].endswith("argv.py") +assert len(sys.argv) == 3 +assert sys.argv[1] == "foobar" +assert sys.argv[2] == "baz" + +out = ~"echo $1 and $2" +assert out == "foobar and baz\n" diff --git a/tests/test_files/injection.py b/tests/test_files/injection.py new file mode 100644 index 0000000..83fd930 --- /dev/null +++ b/tests/test_files/injection.py @@ -0,0 +1,6 @@ +x = ~"uname -p" +print(x in ("arm\n", "x86_64\n")) + +command = "uname -p" +_, _, rc = ~f"{command}" # This doesn't work +print(rc) diff --git a/tests/zxpy_test.py b/tests/zxpy_test.py index 480a777..b371d4f 100644 --- a/tests/zxpy_test.py +++ b/tests/zxpy_test.py @@ -79,7 +79,8 @@ def test_prints(capsys: pytest.CaptureFixture[str]) -> None: def test_argv() -> None: test_file = "./tests/test_files/argv.py" - subprocess.run(["zxpy", test_file]) + returncode = subprocess.check_call(["zxpy", test_file, "--", "foobar", "baz"]) + assert returncode == 0 def test_raise() -> None: @@ -114,3 +115,60 @@ def f(n): assert stderr == b'\n' outlines = [line for line in stdout.decode().splitlines() if line.startswith('>>>')] assert outlines == [">>> hi", ">>> 10", ">>> ... ... ... ... >>> 8", ">>> "] + + +@pytest.mark.parametrize( + ("input", "index", "output"), + ( + ("echo 'hello world' hi", 0, False), + ("echo 'hello world' hi", 4, False), + ("echo 'hello world' hi", 5, True), + ("echo 'hello world' hi", 6, True), + ("echo 'hello world' hi", 16, True), + ("echo 'hello world' hi", 17, True), + ("echo 'hello world' hi", 18, False), + ("echo 'hello world' hi", 21, False), + ('abc "def\'ghi" jkl \'mnop\'', 5, False), + ('abc "def\'ghi" jkl \'mnop\'', 8, False), + ('abc "def\'ghi" jkl \'mnop\'', 10, False), + ('abc "def\'ghi" jkl \'mnop\'', 14, False), + ('abc "def\'ghi" jkl \'mnop\'', 17, False), + ('abc "def\'ghi" jkl \'mnop\'', 18, True), + ('abc "def\'ghi" jkl \'mnop\'', 21, True), + ("'a' 'b' c 'de' 'fg' h", 1, True), + ("'a' 'b' c 'de' 'fg' h", 3, False), + ("'a' 'b' c 'de' 'fg' h", 6, True), + ("'a' 'b' c 'de' 'fg' h", 10, False), + ("'a' 'b' c 'de' 'fg' h", 14, True), + ("'a' 'b' c 'de' 'fg' h", 16, False), + ("'a' 'b' c 'de' 'fg' h", 19, True), + ("'a' 'b' c 'de' 'fg' h", 22, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 1, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 2, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 4, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 6, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 8, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 10, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 12, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 13, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 14, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 15, False), + ("a \"b'c'd'e\" '\"' '\"abc'", 16, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 17, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 18, True), + ("a \"b'c'd'e\" '\"' '\"abc'", 20, True), + ), +) +def test_is_inside_single_quotes(input, index, output) -> None: + assert zx.is_inside_single_quotes(input, index) == output + + +def test_shell_injection(): + """Test injecting commands or shell args like `$0` into shell strings.""" + file = "./tests/test_files/injection.py" + output = subprocess.check_output(["zxpy", file, "--", "abc"]).decode() + assert output == ( + "True\n" # uname -p worked as a string + "127\n" # uname -p inside f-string got quoted + ) + # TODO: $1 injection test diff --git a/zx.py b/zx.py index 7ee2337..cabb4ef 100644 --- a/zx.py +++ b/zx.py @@ -29,6 +29,8 @@ import codecs import contextlib import inspect +import pipes +import re import shlex import subprocess import sys @@ -64,10 +66,21 @@ def cli() -> None: ) parser.add_argument('filename', help='Name of file to run', nargs='?') + # Everything passed after a `--` is arguments to be used by the script itself. + script_args = ['/bin/sh'] + try: + separator_index = sys.argv.index('--') + script_args.extend(sys.argv[separator_index + 1 :]) + # Remove everything after `--` so that argparse passes + sys.argv = sys.argv[:separator_index] + except ValueError: + # `--` not present in command, so no extra script args + pass + args = parser.parse_args(namespace=ZxpyArgs()) - # Remove zxpy executable from argv - del sys.argv[0] + # Once arg parsing is done, replace argv with script args + sys.argv = script_args if args.filename is None: setup_zxpy_repl() @@ -91,9 +104,70 @@ def cli() -> None: install() +def is_inside_single_quotes(string: str, index: int) -> bool: + """Returns True if the given index is inside single quotes in a shell command.""" + quote_index = string.find("'") + if quote_index == -1: + # No single quotes + return False + + if index < quote_index: + # We're before the start of the single quotes + return False + + double_quote_index = string.find('"') + if double_quote_index >= 0 and double_quote_index < quote_index: + next_double_quote = string.find('"', double_quote_index + 1) + if next_double_quote == -1: + # Double quote opened but never closed + return False + + # Single quotes didn't start and we passed the index + if next_double_quote >= index: + return False + + # Ignore all single quotes inside double quotes. + index -= next_double_quote + 1 + rest = string[next_double_quote + 1 :] + return is_inside_single_quotes(rest, index) + + next_quote = string.find("'", quote_index + 1) + if next_quote >= index: + # We're inside single quotes + return True + + index -= next_quote + 1 + rest = string[next_quote + 1 :] + return is_inside_single_quotes(rest, index) + + @contextlib.contextmanager def create_shell_process(command: str) -> Generator[IO[bytes], None, None]: """Creates a shell process, yielding its stdout to read data from.""" + # shell argument support, i.e. $0, $1 etc. + + dollar_indices = [index for index, char in enumerate(command) if char == '$'] + for dollar_index in reversed(dollar_indices): + if ( + dollar_index >= 0 + and dollar_index + 1 < len(command) + and command[dollar_index + 1].isdigit() + and not is_inside_single_quotes(command, dollar_index) + ): + end_index = dollar_index + 1 + while end_index + 1 < len(command) and command[end_index + 1].isdigit(): + end_index += 1 + + number = int(command[dollar_index + 1 : end_index + 1]) + + # Get argument number from sys.argv + if number < len(sys.argv): + replacement = sys.argv[number] + else: + replacement = "" + + command = command[:dollar_index] + replacement + command[end_index + 1 :] + process = subprocess.Popen( command, stdout=subprocess.PIPE,