#!/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 == 59:
                    print("Badgers Removed:", jdata)

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('-b', type=str, required=False, help="Badger IDs to remove seperated by comma with no space", metavar='')
    args = parser.parse_args()
    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
    badgerList = args.b.split(',')
    print("Removing: ", badgerList)
    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': user }, 'task':0}))
        threading.Thread(target=readerThread, args=(ws,)).start()
        apiRequest = json.dumps({
            'bgr_list': badgerList,
            'task': 59
        })
        ws.send(apiRequest)

if __name__ == "__main__":
    main()