1 | from contextlib import contextmanager |
||
2 | import shlex |
||
3 | from subprocess import PIPE, Popen |
||
4 | |||
5 | from coalib.parsing.StringProcessing import escape |
||
6 | |||
7 | |||
8 | @contextmanager |
||
0 ignored issues
–
show
Comprehensibility
Best Practice
introduced
by
![]() |
|||
9 | def run_interactive_shell_command(command, **kwargs): |
||
10 | """ |
||
11 | Runs a single command in shell and provides stdout, stderr and stdin |
||
12 | streams. |
||
13 | |||
14 | This function creates a context manager that sets up the process (using |
||
15 | ``subprocess.Popen()``), returns to caller and waits for process to exit on |
||
16 | leaving. |
||
17 | |||
18 | By default the process is opened in ``universal_newlines`` mode and creates |
||
19 | pipes for all streams (stdout, stderr and stdin) using ``subprocess.PIPE`` |
||
20 | special value. These pipes are closed automatically, so if you want to get |
||
21 | the contents of the streams you should retrieve them before the context |
||
22 | manager exits. |
||
23 | |||
24 | >>> with run_interactive_shell_command(["echo", "TEXT"]) as p: |
||
25 | ... stdout = p.stdout |
||
26 | ... stdout_text = stdout.read() |
||
27 | >>> stdout_text |
||
28 | 'TEXT\\n' |
||
29 | >>> stdout.closed |
||
30 | True |
||
31 | |||
32 | Custom streams provided are not closed except of ``subprocess.PIPE``. |
||
33 | |||
34 | >>> from tempfile import TemporaryFile |
||
35 | >>> stream = TemporaryFile() |
||
36 | >>> with run_interactive_shell_command(["echo", "TEXT"], |
||
37 | ... stdout=stream) as p: |
||
38 | ... stderr = p.stderr |
||
39 | >>> stderr.closed |
||
40 | True |
||
41 | >>> stream.closed |
||
42 | False |
||
43 | |||
44 | :param command: The command to run on shell. This parameter can either |
||
45 | be a sequence of arguments that are directly passed to |
||
46 | the process or a string. A string gets splitted beforehand |
||
47 | using ``shlex.split()``. If providing ``shell=True`` as a |
||
48 | keyword-argument, no ``shlex.split()`` is performed and the |
||
49 | command string goes directly to ``subprocess.Popen()``. |
||
50 | :param kwargs: Additional keyword arguments to pass to |
||
51 | ``subprocess.Popen`` that are used to spawn the process. |
||
52 | :return: A context manager yielding the process started from the |
||
53 | command. |
||
54 | """ |
||
55 | if not kwargs.get("shell", False) and isinstance(command, str): |
||
0 ignored issues
–
show
Comprehensibility
Best Practice
introduced
by
|
|||
56 | command = shlex.split(command) |
||
57 | |||
58 | args = {"stdout": PIPE, |
||
0 ignored issues
–
show
Comprehensibility
Best Practice
introduced
by
|
|||
59 | "stderr": PIPE, |
||
60 | "stdin": PIPE, |
||
61 | "universal_newlines": True} |
||
62 | args.update(kwargs) |
||
0 ignored issues
–
show
Comprehensibility
Best Practice
introduced
by
|
|||
63 | |||
64 | process = Popen(command, **args) |
||
65 | try: |
||
66 | yield process |
||
67 | finally: |
||
68 | if args["stdout"] is PIPE: |
||
69 | process.stdout.close() |
||
70 | if args["stderr"] is PIPE: |
||
71 | process.stderr.close() |
||
72 | if args["stdin"] is PIPE: |
||
73 | process.stdin.close() |
||
74 | |||
75 | process.wait() |
||
76 | |||
77 | |||
78 | def run_shell_command(command, stdin=None, **kwargs): |
||
79 | """ |
||
80 | Runs a single command in shell and returns the read stdout and stderr data. |
||
81 | |||
82 | This function waits for the process (created using ``subprocess.Popen()``) |
||
83 | to exit. Effectively it wraps ``run_interactive_shell_command()`` and uses |
||
84 | ``communicate()`` on the process. |
||
85 | |||
86 | See also ``run_interactive_shell_command()``. |
||
87 | |||
88 | :param command: The command to run on shell. This parameter can either |
||
89 | be a sequence of arguments that are directly passed to |
||
90 | the process or a string. A string gets splitted beforehand |
||
91 | using ``shlex.split()``. |
||
92 | :param stdin: Initial input to send to the process. |
||
93 | :param kwargs: Additional keyword arguments to pass to |
||
94 | ``subprocess.Popen`` that is used to spawn the process. |
||
95 | :return: A tuple with ``(stdoutstring, stderrstring)``. |
||
96 | """ |
||
97 | with run_interactive_shell_command(command, **kwargs) as p: |
||
0 ignored issues
–
show
Comprehensibility
Best Practice
introduced
by
Comprehensibility
Best Practice
introduced
by
Comprehensibility
Best Practice
introduced
by
|
|||
98 | ret = p.communicate(stdin) |
||
99 | return ret |
||
100 | |||
101 | |||
102 | def get_shell_type(): # pragma: no cover |
||
103 | """ |
||
104 | Finds the current shell type based on the outputs of common pre-defined |
||
105 | variables in them. This is useful to identify which sort of escaping |
||
106 | is required for strings. |
||
107 | |||
108 | :return: The shell type. This can be either "powershell" if Windows |
||
109 | Powershell is detected, "cmd" if command prompt is been |
||
110 | detected or "sh" if it's neither of these. |
||
111 | """ |
||
112 | out = run_shell_command("echo $host.name", shell=True)[0] |
||
113 | if out.strip() == "ConsoleHost": |
||
114 | return "powershell" |
||
115 | out = run_shell_command("echo $0", shell=True)[0] |
||
116 | if out.strip() == "$0": |
||
117 | return "cmd" |
||
118 | return "sh" |
||
119 | |||
120 | |||
121 | def prepare_string_argument(string, shell=get_shell_type()): |
||
122 | """ |
||
123 | Prepares a string argument for being passed as a parameter on shell. |
||
124 | |||
125 | On ``sh`` this function effectively encloses the given string |
||
126 | with quotes (either '' or "", depending on content). |
||
127 | |||
128 | :param string: The string to prepare for shell. |
||
129 | :param shell: The shell platform to prepare string argument for. |
||
130 | If it is not "sh" it will be ignored and return the |
||
131 | given string without modification. |
||
132 | :return: The shell-prepared string. |
||
133 | """ |
||
134 | if shell == "sh": |
||
135 | return '"' + escape(string, '"') + '"' |
||
136 | else: |
||
137 | return string |
||
138 | |||
139 | |||
140 | def escape_path_argument(path, shell=get_shell_type()): |
||
141 | """ |
||
142 | Makes a raw path ready for using as parameter in a shell command (escapes |
||
143 | illegal characters, surrounds with quotes etc.). |
||
144 | |||
145 | :param path: The path to make ready for shell. |
||
146 | :param shell: The shell platform to escape the path argument for. Possible |
||
147 | values are "sh", "powershell", and "cmd" (others will be |
||
148 | ignored and return the given path without modification). |
||
149 | :return: The escaped path argument. |
||
150 | """ |
||
151 | if shell == "cmd": |
||
152 | # If a quote (") occurs in path (which is illegal for NTFS file |
||
153 | # systems, but maybe for others), escape it by preceding it with |
||
154 | # a caret (^). |
||
155 | return '"' + escape(path, '"', '^') + '"' |
||
156 | elif shell == "sh": |
||
157 | return escape(path, " ") |
||
158 | else: |
||
159 | # Any other non-supported system doesn't get a path escape. |
||
160 | return path |
||
161 |