Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 0b71c85

Browse files
bdamelebdamele
committed
refactoring, code cleanup, more security-related headers and first /scan method implementation (issue #297)
1 parent a2a71bb commit 0b71c85

1 file changed

Lines changed: 89 additions & 63 deletions

File tree

‎lib/utils/restapi.py‎

Lines changed: 89 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import argparse
99
import os
1010
import sys
11+
import threading
1112

1213
try:
1314
import simplejson as json
@@ -39,135 +40,160 @@
3940
from lib.core.settings import UNICODE_ENCODING
4041
from lib.core.settings import RESTAPI_SERVER_PORT
4142

42-
43-
# local global variables
44-
session_ids = []
45-
admin_id = ""
46-
43+
# Local global variables
44+
options = AttribDict()
45+
adminid = ""
46+
tasks = []
4747

4848
# Generic functions
4949
def jsonize(data):
50-
#return json.dumps(data, sort_keys=False, indent=4)
5150
return json.dumps(data, sort_keys=False)
5251

53-
54-
def is_admin(session_id):
55-
global admin_id
56-
#print "[INFO] Admin ID: %s" % admin_id
57-
#print "[INFO] Session ID: %s" % session_id
58-
if admin_id != session_id:
52+
def is_admin(taskid):
53+
global adminid
54+
#print "[INFO] Admin ID: %s" % adminid
55+
#print "[INFO] Task ID: %s" % taskid
56+
if adminid != taskid:
5957
return False
6058
else:
6159
return True
6260

63-
6461
@hook('after_request')
6562
def security_headers():
6663
"""
6764
Set some headers across all HTTP responses
6865
"""
6966
response.headers["Server"] = "Server"
70-
response.headers["X-Frame-Options"] = "sameorigin"
67+
response.headers["X-Content-Type-Options"] = "nosniff"
68+
response.headers["X-Frame-Options"] = "DENY"
7169
response.headers["X-XSS-Protection"] = "1; mode=block"
70+
response.headers["Pragma"] = "no-cache"
71+
response.headers["Cache-Control"] = "no-cache"
72+
response.headers["Expires"] = "0"
73+
response.content_type = "application/json; charset=UTF-8"
7274

75+
##############################
76+
# HTTP Status Code functions #
77+
##############################
7378

74-
# HTTP Status Code functions
7579
@error(401) # Access Denied
76-
def error401(error):
80+
def error401(error=None):
7781
return "Access denied"
7882

79-
8083
@error(404) # Not Found
81-
def error404(error):
84+
def error404(error=None):
8285
return "Nothing here"
8386

84-
8587
@error(405) # Method Not Allowed (e.g. when requesting a POST method via GET)
86-
def error405(error):
88+
def error405(error=None):
8789
return "Method not allowed"
8890

89-
9091
@error(500) # Internal Server Error
91-
def error500(error):
92+
def error500(error=None):
9293
return "Internal server error"
9394

94-
95-
################################
96-
# Session management functions #
97-
################################
95+
#############################
96+
# Task management functions #
97+
#############################
9898

9999
# Users' methods
100-
@get("/session/new")
101-
def session_new():
100+
@get("/task/new")
101+
def task_new():
102102
"""
103-
Create new session token
103+
Create new task ID
104104
"""
105-
global session_ids
106-
session_id = hexencode(os.urandom(32))
107-
session_ids.append(session_id)
108-
response.content_type = "application/json; charset=UTF-8"
109-
return jsonize({"sessionid": session_id})
110-
105+
global tasks
106+
taskid = hexencode(os.urandom(32))
107+
tasks.append(taskid)
108+
return jsonize({"taskid": taskid})
111109

112-
@post("/session/destroy")
113-
def session_destroy():
110+
@get("/task/<taskid>/destroy")
111+
def task_destroy(taskid):
114112
"""
115-
Destroy own session token
113+
Destroy own task ID
116114
"""
117-
session_id = request.json.get("sessionid", "")
118-
if session_id in session_ids:
119-
session_ids.remove(session_id)
115+
if taskid in tasks:
116+
tasks.remove(taskid)
120117
return jsonize({"success": True})
121118
else:
122-
abort(500)
119+
abort(500, "Invalid task ID")
123120

124121
# Admin's methods
125-
@post("/session/list")
126-
def session_list():
122+
@get("/task/<taskid>/list")
123+
def task_list(taskid):
127124
"""
128-
List all active sessions
125+
List all active tasks
129126
"""
130-
if is_admin(request.json.get("sessionid", "")):
131-
response.content_type = "application/json; charset=UTF-8"
132-
return jsonize({"sessions": session_ids})
127+
if is_admin(taskid):
128+
return jsonize({"tasks": tasks})
133129
else:
134130
abort(401)
135131

136-
137-
@post("/session/flush")
138-
def session_flush():
132+
@get("/task/<taskid>/flush")
133+
def task_flush(taskid):
139134
"""
140-
Flush session spool (destroy all sessions)
135+
Flush task spool (destroy all tasks)
141136
"""
142-
global session_ids
143-
if is_admin(request.json.get("sessionid", "")):
144-
session_ids = []
137+
global tasks
138+
if is_admin(taskid):
139+
tasks = []
145140
return jsonize({"success": True})
146141
else:
147142
abort(401)
148143

144+
##################################
145+
# sqlmap core interact functions #
146+
##################################
147+
@post("/scan/<taskid>")
148+
def scan(taskid):
149+
"""
150+
Mount a scan with sqlmap
151+
"""
152+
global options
153+
154+
if taskid not in tasks:
155+
abort(500, "Invalid task ID")
149156

150-
@post("/download/<target>/<filename:path>")
151-
def download(target, filename):
157+
# Initialize sqlmap engine's options with user's provided options
158+
# within the JSON request
159+
for key, value in request.json.items():
160+
if key != "taskid":
161+
options[key] = value
162+
init(options, True)
163+
164+
# Launch sqlmap engine in a separate thread
165+
thread = threading.Thread(target=start)
166+
thread.daemon = True
167+
thread.start()
168+
169+
return jsonize({"success": True})
170+
171+
@post("/download/<taskid>/<target>/<filename:path>")
172+
def download(taskid, target, filename):
152173
"""
153174
Download a certain file from the file system
154175
"""
176+
if taskid not in tasks:
177+
abort(500, "Invalid task ID")
178+
155179
path = os.path.join(paths.SQLMAP_OUTPUT_PATH, target)
156180
if os.path.exists(path):
157181
return static_file(filename, root=path)
158182
else:
159183
abort(500)
160184

161-
162185
def restAPIrun(host="0.0.0.0", port=RESTAPI_SERVER_PORT):
163186
"""
164187
Initiate REST-JSON API
165188
"""
166-
global admin_id
167-
admin_id = hexencode(os.urandom(32))
189+
global adminid
190+
global options
191+
global tasks
192+
adminid = hexencode(os.urandom(32))
193+
tasks.append(adminid)
168194
options = AttribDict(cmdLineOptions)
169195
logger.info("Running REST-JSON API server at '%s:%d'.." % (host, port))
170-
logger.info("The admin session ID is: %s" % admin_id)
196+
logger.info("The admin task ID is: %s" % adminid)
171197
run(host=host, port=port)
172198

173199
def client(host, port):
@@ -176,8 +202,8 @@ def client(host, port):
176202

177203
# TODO: write a simple client with urllib2, for now use curl from command line
178204
print "[ERROR] Not yet implemented, use curl from command line instead for now, for example:"
179-
print "\n\t$ curl --proxy http://127.0.0.1:8080 http://%s:%s/session/new" % (host, port)
180-
print "\t$ curl --proxy http://127.0.0.1:8080 -H \"Content-Type: application/json\" -X POST -d '{\"sessionid\": \"<admin session id>\"}' http://%s:%d/session/list\n" % (host, port)
205+
print "\n\t$ curl --proxy http://127.0.0.1:8080 http://127.0.0.1:%s/task/new" % port
206+
print "\t$ curl --proxy http://127.0.0.1:8080 -H \"Content-Type: application/json\" -X POST -d '{\"targetUrl\": \"<target URL>\"}' http://127.0.0.1:%d/scan/<task ID>\n" % port
181207

182208
if __name__ == "__main__":
183209
"""

0 commit comments

Comments
 (0)