<commit_before>#!/usr/bin/env python3

import json
import os
import shutil
import sqlite3
import sys
import threading

from collections import OrderedDict, defaultdict
from flask import Flask, make_response, send_file
from flask_restful import reqparse, abort, Api, Resource
from time import gmtime, strftime
from locale import getlocale
from subprocess import Popen, PIPE

# TODO Fix exports
app = Flask(__name__, static_url_path=os.path.dirname(os.path.abspath(__file__)))
api = Api(app)

parser = reqparse.RequestParser()
# for /tools
parser.add_argument('search')
# for /jobs
parser.add_argument('id')
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')

class RunThreads (threading.Thread):
    def __init__(self, tool, job):
        threading.Thread.__init__(self)
        self.tool_stdout = ''
        self.tool_sterr = ''
        self.tool = tool
        self.job = job

    def execute_tool(self, job):
        # Always check any tools provided by
        # community members
        print("Running")
        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
        zip_file = None
        if return_code == 0:
            zip_file = job['target_name'] + '_' + str(job['id'])
            shutil.make_archive(zip_file, 'zip', 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()
        cur.execute("UPDATE 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)

class Pong(Resource):
    def get(self):
        return { 'message':'pong' }

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 Jobs(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 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 jobs")
        data = dict(result=[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=<PASSWORD>&user_file=/user/file&password_file=/password/file" http://127.0.0.1:5000/jobs
        cur.execute("INSERT INTO jobs(tool,target,target_name,protocol,port_number,user,password,user_file,password_file) VALUES(?,?,?,?,?,?,?,?,?)",(args['tool'],args['target'],args['target_name'],args['protocol'],args['port_number'],args['user'],args['password'],args['user_file'],args['password_file']))
        data = cur.lastrowid
        con.commit()

        # Close connection
        if con:
            con.close()

        return { 'id':data }, 201

# List individual jobs
class JobsId(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 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 JobsIdExecute(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 jobs WHERE id = ?",(args['id'],))
        job_data = dict(result=[dict(r) for r in cur.fetchall()])

        # Index is now tied to database schema, yuck
        print(args['id'])
        job = job_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.getcwd() + '/' + strftime("%Y%m%d") + \
                                "_autopwn_" + \
                                job_data['result'][0]['target_name']
        os.makedirs(job['output_dir'])

        tool['id'] = job['tool']
        # Get dependencies
        cur.execute("SELECT dependency from dependencies WHERE tool = ?",(tool['id'],))
        dependency = dict(result=[dict(r) for r in cur.fetchall()])

        # Get tool execute string
        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]

        # 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)
        print(tool['execute_string'])

        thread = []
        # Tool string generated, execute
        thread.append(RunThreads(tool,job))
        # If main process dies, everything else *SHOULD* as well
        thread[-1].daemon = True
        # Start threads
        thread[-1].start()

        return {'message':'Job Id to be executed here'}, 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 Exports(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 jobs")
        data = dict(result=[dict(r) for r in cur.fetchall()])
        print(data)

        # Close connection
        if con:
            con.close()

        return data

# Retrieve output for job id
class ExportsId(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 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

# 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 all jobs
# curl -i http://127.0.0.1:5000/jobs
api.add_resource(Jobs, '/jobs')
# Fetch job id
# curl -i http://127.0.0.1:5000/jobs/1
api.add_resource(JobsId, '/jobs/<job_id>')
# Execute job id
# curl -i --data "id=1" http://127.0.0.1:5000/jobs/execute
api.add_resource(JobsIdExecute, '/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>')
# 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(Exports, '/exports')
# Fetch tool output from job id
# curl -i http://127.0.0.1:5000/exports/1
api.add_resource(ExportsId, '/exports/<job_id>')

if __name__ == '__main__':
    print(os.path.dirname(os.path.abspath(__file__)))
    app.run(debug=True,threaded=True)

def main():
    if os.path.isfile('/.dockerinit'):
        print("Running in docker")
        app.run(host='0.0.0.0', debug=True,threaded=True)
    else:
        app.run(debug=True,threaded=True)
<commit_msg>Export path bug fix<commit_after>#!/usr/bin/env python3

import json
import os
import shutil
import sqlite3
import sys
import threading

from collections import OrderedDict, defaultdict
from flask import Flask, make_response, send_file
from flask_restful import reqparse, abort, Api, Resource
from time import gmtime, strftime
from locale import getlocale
from subprocess import Popen, PIPE

# TODO Fix exports
app = Flask(__name__, static_url_path=os.path.dirname(os.path.abspath(__file__)))
api = Api(app)

parser = reqparse.RequestParser()
# for /tools
parser.add_argument('search')
# for /jobs
parser.add_argument('id')
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')

class RunThreads (threading.Thread):
    def __init__(self, tool, job):
        threading.Thread.__init__(self)
        self.tool_stdout = ''
        self.tool_sterr = ''
        self.tool = tool
        self.job = job

    def execute_tool(self, job):
        # Always check any tools provided by
        # community members
        print("Running")
        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
        zip_file = None
        if return_code == 0:
            zip_file = job['target_name'] + '_' + str(job['id'])
            shutil.make_archive(zip_file, 'zip', 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()
        cur.execute("UPDATE 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)

class Pong(Resource):
    def get(self):
        return { 'message':'pong' }

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 Jobs(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 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 jobs")
        data = dict(result=[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 jobs(tool,target,target_name,protocol,port_number,user,password,user_file,password_file) VALUES(?,?,?,?,?,?,?,?,?)",(args['tool'],args['target'],args['target_name'],args['protocol'],args['port_number'],args['user'],args['password'],args['user_file'],args['password_file']))
        data = cur.lastrowid
        con.commit()

        # Close connection
        if con:
            con.close()

        return { 'id':data }, 201

# List individual jobs
class JobsId(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 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 JobsIdExecute(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 jobs WHERE id = ?",(args['id'],))
        job_data = dict(result=[dict(r) for r in cur.fetchall()])

        # Index is now tied to database schema, yuck
        print(args['id'])
        job = job_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") + \
                                "_autopwn_" + \
                                job_data['result'][0]['target_name']
        os.makedirs(job['output_dir'])

        tool['id'] = job['tool']
        # Get dependencies
        cur.execute("SELECT dependency from dependencies WHERE tool = ?",(tool['id'],))
        dependency = dict(result=[dict(r) for r in cur.fetchall()])

        # Get tool execute string
        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]

        # 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)
        print(tool['execute_string'])

        thread = []
        # Tool string generated, execute
        thread.append(RunThreads(tool,job))
        # If main process dies, everything else *SHOULD* as well
        thread[-1].daemon = True
        # Start threads
        thread[-1].start()

        return {'message':'Job Id to be executed here'}, 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 Exports(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 jobs")
        data = dict(result=[dict(r) for r in cur.fetchall()])
        print(data)

        # Close connection
        if con:
            con.close()

        return data

# Retrieve output for job id
class ExportsId(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 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

# 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 all jobs
# curl -i http://127.0.0.1:5000/jobs
api.add_resource(Jobs, '/jobs')
# Fetch job id
# curl -i http://127.0.0.1:5000/jobs/1
api.add_resource(JobsId, '/jobs/<job_id>')
# Execute job id
# curl -i --data "id=1" http://127.0.0.1:5000/jobs/execute
api.add_resource(JobsIdExecute, '/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>')
# 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(Exports, '/exports')
# Fetch tool output from job id
# curl -i http://127.0.0.1:5000/exports/1
api.add_resource(ExportsId, '/exports/<job_id>')

if __name__ == '__main__':
    print(os.path.dirname(os.path.abspath(__file__)))
    app.run(debug=True,threaded=True)

def main():
    if os.path.isfile('/.dockerinit'):
        print("Running in docker")
        app.run(host='0.0.0.0', debug=True,threaded=True)
    else:
        app.run(debug=True,threaded=True)

{'fraction_non_alphanumeric': 0.14854315088946718, 'fraction_numerical': 0.013821935924235314, 'mean_word_length': 2.695727573703295, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 24, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2181134', 'n_tokens_mistral': 7648, 'n_tokens_neox': 7270, 'n_words': 1655}