1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
|
"""
The script object representing mitmproxy inline scripts.
Script objects know nothing about mitmproxy or mitmproxy's API - this knowledge is provided
by the mitmproxy-specific ScriptContext.
"""
# Do not import __future__ here, this would apply transitively to the inline scripts.
from __future__ import absolute_import, print_function, division
import inspect
import os
import shlex
import sys
import contextlib
import warnings
import six
from mitmproxy import exceptions
@contextlib.contextmanager
def setargs(args):
oldargs = sys.argv
sys.argv = args
try:
yield
finally:
sys.argv = oldargs
class Script(object):
"""
Script object representing an inline script.
"""
def __init__(self, command, context):
self.command = command
self.args = self.parse_command(command)
self.ctx = context
self.ns = None
def __enter__(self):
self.load()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_val:
return False # reraise the exception
self.unload()
@property
def filename(self):
return self.args[0]
@staticmethod
def parse_command(command):
if not command or not command.strip():
raise exceptions.ScriptException("Empty script command.")
# Windows: escape all backslashes in the path.
if os.name == "nt": # pragma: no cover
backslashes = shlex.split(command, posix=False)[0].count("\\")
command = command.replace("\\", "\\\\", backslashes)
args = shlex.split(command) # pragma: no cover
args[0] = os.path.expanduser(args[0])
if not os.path.exists(args[0]):
raise exceptions.ScriptException(
("Script file not found: %s.\r\n"
"If your script path contains spaces, "
"make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") %
args[0])
elif os.path.isdir(args[0]):
raise exceptions.ScriptException("Not a file: %s" % args[0])
return args
def load(self):
"""
Loads an inline script.
Returns:
The return value of self.run("start", ...)
Raises:
ScriptException on failure
"""
if self.ns is not None:
raise exceptions.ScriptException("Script is already loaded")
script_dir = os.path.dirname(os.path.abspath(self.args[0]))
self.ns = {'__file__': os.path.abspath(self.args[0])}
sys.path.append(script_dir)
sys.path.append(os.path.join(script_dir, ".."))
try:
with open(self.filename) as f:
code = compile(f.read(), self.filename, 'exec')
exec(code, self.ns, self.ns)
except Exception:
six.reraise(
exceptions.ScriptException,
exceptions.ScriptException.from_exception_context(),
sys.exc_info()[2]
)
finally:
sys.path.pop()
sys.path.pop()
start_fn = self.ns.get("start")
if start_fn and len(inspect.getargspec(start_fn).args) == 2:
warnings.warn(
"The 'args' argument of the start() script hook is deprecated. "
"Please use sys.argv instead."
)
return self.run("start", self.args)
return self.run("start")
def unload(self):
try:
return self.run("done")
finally:
self.ns = None
def run(self, name, *args, **kwargs):
"""
Runs an inline script hook.
Returns:
The return value of the method.
None, if the script does not provide the method.
Raises:
ScriptException if there was an exception.
"""
if self.ns is None:
raise exceptions.ScriptException("Script not loaded.")
f = self.ns.get(name)
if f:
try:
with setargs(self.args):
return f(self.ctx, *args, **kwargs)
except Exception:
six.reraise(
exceptions.ScriptException,
exceptions.ScriptException.from_exception_context(),
sys.exc_info()[2]
)
else:
return None
|