forked from Gallopsled/pwntools
-
Notifications
You must be signed in to change notification settings - Fork 0
/
internal.py
163 lines (129 loc) · 4.61 KB
/
internal.py
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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
from __future__ import absolute_import
from __future__ import division
import os
import sys
from pwnlib.context import context
__all__ = ['make_function']
loaded = {}
lookup = None
def init_mako():
global lookup, render_global
from mako.lookup import TemplateLookup
from mako.parsetree import Tag, Text
from mako import ast
import threading
if lookup is not None:
return
class IsInsideManager(object):
def __init__(self, parent):
self.parent = parent
def __enter__(self):
self.oldval = self.parent.is_inside
self.parent.is_inside = True
return self.oldval
def __exit__(self, *args):
self.parent.is_inside = self.oldval
class IsInside(threading.local):
is_inside = False
def go_inside(self):
return IsInsideManager(self)
render_global = IsInside()
cache = context.cache_dir
if cache:
cache = os.path.join(cache, 'mako')
curdir = os.path.dirname(os.path.abspath(__file__))
lookup = TemplateLookup(
directories = [os.path.join(curdir, 'templates')],
module_directory = cache
)
# The purpose of this definition is to create a new Tag.
# The Tag has a metaclass, which saves this definition even
# though to do not use it here.
class pwn_docstring(Tag):
__keyword__ = 'docstring'
def __init__(self, *args, **kwargs):
super(pwn_docstring, self).__init__('docstring', (), (), (), (), **kwargs)
self.ismodule = True
@property
def text(self):
children = self.get_children()
if len(children) != 1 or not isinstance(children[0], Text):
raise Exception("docstring tag only supports text")
docstring = children[0].content
return '__doc__ = %r' % docstring
@property
def code(self):
return ast.PythonCode(self.text)
def accept_visitor(self, visitor):
method = getattr(visitor, "visitCode", lambda x: x)
method(self)
def lookup_template(filename):
init_mako()
if filename not in loaded:
loaded[filename] = lookup.get_template(filename)
return loaded[filename]
def get_context_from_dirpath(directory):
"""
>>> get_context_from_dirpath('common')
{}
>>> get_context_from_dirpath('i386')
{'arch': 'i386'}
>>> get_context_from_dirpath('amd64/linux') == {'arch': 'amd64', 'os': 'linux'}
True
"""
parts = directory.split(os.path.sep)
arch = osys = None
if len(parts) > 0:
arch = parts[0]
if len(parts) > 1:
osys = parts[1]
if osys == 'common':
osys = None
if arch == 'common':
arch = None
return {'os': osys, 'arch': arch}
def make_function(funcname, filename, directory):
import functools
import inspect
path = os.path.join(directory, filename)
template = lookup_template(path)
local_ctx = get_context_from_dirpath(directory)
def res(*args, **kwargs):
with render_global.go_inside() as was_inside:
with context.local(**local_ctx):
lines = template.render(*args, **kwargs).split('\n')
for i, line in enumerate(lines):
def islabelchar(c):
return c.isalnum() or c == '.' or c == '_'
if ':' in line and islabelchar(line.lstrip()[0]):
line = line.lstrip()
elif line.startswith(' '):
line = ' ' + line.lstrip()
lines[i] = line
while lines and not lines[-1]: lines.pop()
while lines and not lines[0]: lines.pop(0)
s = '\n'.join(lines)
while '\n\n\n' in s:
s = s.replace('\n\n\n', '\n\n')
if was_inside:
return s
else:
return s + '\n'
# Setting _relpath is a slight hack only used to get better documentation
res._relpath = path
res.__module__ = 'pwnlib.shellcraft.' + os.path.dirname(path).replace('/','.')
res.__name__ = res.__qualname__ = funcname
res.__doc__ = inspect.cleandoc(template.module.__doc__ or '')
if hasattr(inspect, 'signature'):
sig = inspect.signature(template.module.render_body)
sig = sig.replace(parameters=list(sig.parameters.values())[1:-1])
res.__signature__ = sig
@functools.wraps(res)
def function(*a):
return sys.modules[res.__module__].function(res.__name__, res, *a)
@functools.wraps(res)
def call(*a):
return sys.modules[res.__module__].call(res.__name__, *a)
res.function = function
res.call = call
return res