#!/usr/bin/env python3
import errno
import json
import locale
import os
import shutil
#import ssl
import sqlite3
import sys
import threading
from collections import OrderedDict, defaultdict
from flask import Flask, make_response, send_file, jsonify
from flask.ext.cors import CORS
from flask_restful import reqparse, abort, Api, Resource
from time import gmtime, strftime
from locale import getlocale
from subprocess import Popen, PIPE
# TODO
# - TLS
app = Flask(__name__, static_url_path=os.path.dirname(os.path.abspath(__file__)))
CORS(app)
api = Api(app)
parser = reqparse.RequestParser()
# for /tools
parser.add_argument('search')
# for /assessments/jobs and /tools/jobs
parser.add_argument('id')
parser.add_argument('assessment')
parser.add_argument('tool')
parser.add_argument('target')
parser.add_argument('target_name')
parser.add_argument('protocol')
parser.add_argument('port_number')
parser.add_argument('user')
parser.add_argument('password')
parser.add_argument('user_file')
parser.add_argument('password_file')
parser.add_argument('url')
class RunThreads (threading.Thread):
def __init__(self, tool, job, context):
threading.Thread.__init__(self)
self.tool_stdout = ''
self.tool_sterr = ''
self.tool = tool
self.job = job
self.context = context
def execute_tool(self, job, context):
# Always check any tools provided by
# community members
print("Running")
print("Execute string " + self.tool['execute_string'])
proc = Popen(self.tool['execute_string'], stdout=PIPE, stderr=PIPE, shell=True)
decode_locale = lambda s: s.decode(getlocale()[1])
self.tool_stdout, self.tool_stderr = map(decode_locale, proc.communicate())
# Callback / pause from here
return_code = proc.returncode
# Zip resulting directory - This is crap because when doing assessment
# this will be run at each tool execution.
zip_file = os.path.dirname(os.path.abspath(__file__)) + \
"/" + job['target_name'] + '_' + str(job['id'])
shutil.make_archive(zip_file, 'zip', job['output_dir'])
# TODO This could be dodgy
shutil.rmtree(job['output_dir'])
# Update completed and return_code field in db
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
cur = con.cursor()
if context == 'tool':
cur.execute("UPDATE tool_jobs SET executed = 1, return_code = ?, zip_file = ? WHERE id = ?",(str(return_code),str(zip_file),str(job['id'])))
con.commit()
if context == 'assessment':
# Pull and check for 0
cur.execute("SELECT return_code FROM assessment_jobs WHERE id = ? AND return_code == 0",(str(job['id'])))
data = cur.fetchall()
if len(data) == 0 or data[0][0] == 0:
cur.execute("UPDATE assessment_jobs SET executed = 1, return_code = ?, zip_file = ? WHERE id = ?",(str(return_code),str(zip_file),str(job['id'])))
con.commit()
# Close connection
if con:
con.close()
def run(self):
self.execute_tool(self.job, self.context)
class Pong(Resource):
def get(self):
return { 'message':'pong' }
class Assessments(Resource):
# Get assessments
def get(self):
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
# If /assessments?search=xxx not specified then SELECT *
if args['search'] != None:
cur.execute("SELECT * FROM assessments WHERE name LIKE ? OR description LIKE ?",('%' + args['search'] + '%','%' + args['search'] + '%'))
else:
cur.execute("SELECT * FROM assessments")
data = dict(result=[dict(r) for r in cur.fetchall()])
# Add tools to assessment
for i, assessment in enumerate(data['result']):
cur.execute("SELECT tool FROM assessment_tools WHERE assessment = ?",(str(i+1),))
tool_ids = dict(result=[dict(r) for r in cur.fetchall()])
assessment['tools'] = tool_ids['result']
# Close connection
if con:
con.close()
return data
# Submit a new tool
def post(self):
args = parser.parse_args()
return args
class AssessmentsId(Resource):
# Get assessments
def get(self, assessment_id):
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT * FROM assessments WHERE id = ?",(assessment_id),)
data = dict(result=[dict(r) for r in cur.fetchall()])
# Add tools to assessment
cur.execute("SELECT tool FROM assessment_tools WHERE assessment = ?",(str(assessment_id),))
tool_ids = dict(result=[dict(r) for r in cur.fetchall()])
data['result'][0]['tools'] = tool_ids['result']
# Close connection
if con:
con.close()
return data
class Tools(Resource):
# Get tools
def get(self):
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
# If /tools?search=xxx not specified then SELECT *
if args['search'] != None:
cur.execute("SELECT * FROM tools WHERE name LIKE ? OR description LIKE ?",('%' + args['search'] + '%','%' + args['search'] + '%'))
else:
cur.execute("SELECT * FROM tools")
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
# Submit a new tool
def post(self):
args = parser.parse_args()
return args
class ToolsId(Resource):
# Get tools
def get(self, tool_id):
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT * FROM tools WHERE id =?",(tool_id,))
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
# Submit a new tool
def post(self):
args = parser.parse_args()
return args
class ToolsJobs(Resource):
# List jobs
def get (self):
#args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
# If /jobs?search=xxx not specified then SELECT *
#if args['search'] != None:
# cur.execute("SELECT * FROM tool_jobs WHERE tool LIKE ? OR target LIKE ? OR target_name LIKE ? OR protocol LIKE ? OR port_number LIKE ? OR user like ? OR password LIKE ? OR user_file LIKE ? OR password_file LIKE ?",('%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%','%' + args['search'] + '%'))
#else:
cur.execute("SELECT * FROM tool_jobs")
data = [dict(r) for r in cur.fetchall()]
# Close connection
if con:
con.close()
return data
# Submit new job
def post(self):
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
cur = con.cursor()
# curl -i --data "tool=1&target=localhost&target_name=target_name&protocol=https&port_number=1337&user=a_user&password=a_password&user_file=/user/file&password_file=/password/file" http://127.0.0.1:5000/jobs
cur.execute("INSERT INTO tool_jobs(tool,target,target_name,protocol,port_number,user,password,user_file,password_file,url) VALUES(?,?,?,?,?,?,?,?,?,?)",(args['tool'],args['target'],args['target_name'].replace('/','_'),args['protocol'],args['port_number'],args['user'],args['password'],args['user_file'],args['password_file'],args['url']))
data = cur.lastrowid
con.commit()
# Close connection
if con:
con.close()
return { 'id':data }, 201
# List individual jobs
class ToolsJobsId(Resource):
def get(self, job_id):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT * FROM tool_jobs WHERE id = ?",(job_id,))
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
# Execute job
class ToolsJobsIdExecute(Resource):
def post(self):
# Process placeholders
tool = {}
job = {}
args = parser.parse_args()
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
# Get job id columns
cur.execute("SELECT * FROM tool_jobs WHERE id = ?",(args['id'],))
job_data = dict(result=[dict(r) for r in cur.fetchall()])
# Index is now tied to database schema, yuck
job = job_data['result'][0]
tool['id'] = job['tool']
# Get tool execute string and tool name
cur.execute("SELECT * FROM tools WHERE id = ?",(tool['id'],))
tool_data = dict(result=[dict(r) for r in cur.fetchall()])
tool = tool_data['result'][0]
# TODO Allow set in config file
job['tools_directory'] = "/root/tools"
job['date'] = strftime("%Y%m%d_%H%M%S%z")
job['output_dir'] = os.path.dirname(os.path.abspath(__file__)) + \
'/' + strftime("%Y%m%d") + \
"_torc_" + \
job_data['result'][0]['target_name'] + \
"_" + job_data['result'][0]['target'].replace('/','_') + \
"_" + str(job_data['result'][0]['id']) + \
"_" + tool['name']
try:
os.makedirs(job['output_dir'])
except OSError as e:
pass
#if e.errno == errno.EEXIST:
# return {'message':'Directory exists'}, 500
# Get dependencies
cur.execute("SELECT dependency from dependencies WHERE tool = ?",(tool['id'],))
dependency = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
ddict_options = defaultdict(lambda:'')
for option in job:
ddict_options[option] = job[option]
tool['execute_string'] = tool['execute_string'].format(**ddict_options)
thread = []
# Tool string generated, execute
thread.append(RunThreads(tool,job,'tool'))
# If main process dies, everything else *SHOULD* as well
thread[-1].daemon = True
# Start threads
thread[-1].start()
return {'message':'Tool executed'}, 201
class Dependencies(Resource):
def get(self):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT * FROM dependency_names")
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
class Options(Resource):
def get(self):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT * FROM options")
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
# Retrieve options for tool
class OptionsId(Resource):
def get(self, tool_id):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT option, required FROM tool_options WHERE tool = ?",(tool_id,))
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
class ToolsExports(Resource):
def get(self):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT id, zip_file FROM tool_jobs")
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
# Retrieve output for job id
class ToolsExportsId(Resource):
def get(self, job_id):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT zip_file FROM tool_jobs WHERE id = ?",(job_id,))
data = dict(result=[dict(r) for r in cur.fetchall()])
zip_file = data['result'][0]['zip_file'] + '.zip'
# Close connection
if con:
con.close()
return send_file(zip_file, as_attachment=True)
# Retrieve dependencies for tool
class DependenciesId(Resource):
def get(self, tool_id):
con = sqlite3.connect(os.path.dirname(os.path.abspath(__file__)) + '/assets.db')
con.row_factory = sqlite3.Row
cur = con.cursor()
cur.execute("SELECT dependency FROM dependencies WHERE tool = ?",(tool_id,))
data = dict(result=[dict(r) for r in cur.fetchall()])
# Close connection
if con:
con.close()
return data
class CheckCertificate():
def __init__(self):
if os.path.isfile(os.path.dirname(os.path.abspath(__file__)) + '/keys/cert.crt') != True:
print("First run, generating certificate..")
execute_string = "openssl req -nodes -x509 -newkey rsa:4096 -keyout " + \
os.path.dirname(os.path.abspath(__file__)) + \
"/keys/key.pem -out " + os.path.dirname(os.path.abspath(__file__)) + \
"/keys/cert.crt -days 356 -subj " + \
"\"/C=AP/ST=Land/L=torc/O=AP Team/subjectAltName=IP:127.0.0.1\""
proc = Popen(execute_string, stdout=PIPE, stderr=PIPE, shell=True)
decode_locale = lambda s: s.decode(getlocale()[1])
self.tool_stdout, self.tool_stderr = map(decode_locale, proc.communicate())
# Callback / pause from here
return_code = proc.returncode
# Pong!
# curl -i http://127.0.0.1:5000/ping
api.add_resource(Pong, '/ping')
# Fetch all tools
# curl -i http://127.0.0.1:5000/tools
api.add_resource(Tools, '/tools')
# Fetch tool id
# curl -i http://127.0.0.1:5000/tools/1
api.add_resource(ToolsId, '/tools/<tool_id>')
# Fetch all assessments
# curl -i http://127.0.0.1:5000/assessments
api.add_resource(Assessments, '/assessments')
# Fetch assessment id
# curl -i http://127.0.0.1:5000/assessments/1
api.add_resource(AssessmentsId, '/assessments/<assessment_id>')
# Fetch all tool jobs
# curl -i http://127.0.0.1:5000/jobs
api.add_resource(ToolsJobs, '/tools/jobs')
# Fetch tool job id
# curl -i http://127.0.0.1:5000/jobs/1
api.add_resource(ToolsJobsId, '/tools/jobs/<job_id>')
# Execute tool job id
# curl -i --data "id=1" http://127.0.0.1:5000/jobs/execute
api.add_resource(ToolsJobsIdExecute, '/tools/jobs/execute')
# Fetch all dependencies
# curl -i http://127.0.0.1:5000/dependencies
api.add_resource(Dependencies, '/dependencies')
# Fetch all dependencies for tool id
api.add_resource(DependenciesId, '/dependencies/<tool_id>')
# TODO Review (/tools/options? /assessment/options?)
# Fetch all options
# curl -i http://127.0.0.1:5000/options
api.add_resource(Options, '/options')
# Fetch all options for tool id
# curl -i http://127.0.0.1:5000/options/1
api.add_resource(OptionsId, '/options/<tool_id>')
# Fetch all tool outputs
# curl -i http://127.0.0.1:5000/exports
api.add_resource(ToolsExports, '/tools/jobs/exports')
# Fetch tool output from job id
# curl -i http://127.0.0.1:5000/exports/1
api.add_resource(ToolsExportsId, '/tools/jobs/exports/<job_id>')
def main():
# TODO Hack, fix
#locale.setlocale(locale.LC_CTYPE, 'en_US.UTF-8')
print(os.path.dirname(os.path.abspath(__file__)))
# Check for certificate (or create if non-existent)
#CheckCertificate()
#context = ssl.SSLContext(ssl.PROTOCOL_TLSv1_2)
#context.load_cert_chain(os.path.dirname(os.path.abspath(__file__)) + "/keys/cert.crt", \
# os.path.dirname(os.path.abspath(__file__)) + "/keys/key.pem")
if os.path.isfile('/.dockerenv'):
print("Running in docker")
#app.run(host='0.0.0.0',debug=True,threaded=True,ssl_context=context)
app.run(host='0.0.0.0',debug=True,threaded=True)
else:
#app.run(debug=True,threaded=True,ssl_context=context)
app.run(debug=True,threaded=True,port=5000)
if __name__ == '__main__':
main()
