sleep_python_bridge/beacongrapher.py at main · 0xTriboulet/sleep_python_bridge · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
#!/usr/local/bin/python3
from sleep_python_bridge.striker import CSConnector

from argparse import ArgumentParser
from pprint import pp, pprint
from pathlib import Path
import json

####################
## Variables
# Initialize lists
beaconsresult = [] # raw data from CS
beaconLogs = [] # Cleaned list to feed final JSON

# JSON file
datafile = "beacons.json"

####################
## FUNCTIONS
def parseArguments():
    parser = ArgumentParser()
    parser.add_argument('host', help='The teamserver host.')
    parser.add_argument('port', help='The teamserver port.')
    parser.add_argument('username', help='The desired username.')
    parser.add_argument('password', help='The teamserver password.')
    parser.add_argument('path', help="Directory to CobaltStrike")

    args = parser.parse_args()
    return args

def main(args):

    cs_host = args.host
    cs_port = args.port
    cs_user = args.username
    cs_pass = args.password
    cs_directory = args.path

    ####################
    ## Connect to server
    print(f"[*] Connecting to teamserver: {cs_host}")
    with CSConnector(
        cs_host=cs_host,
        cs_port=cs_port,
        cs_user=cs_user,
        cs_pass=cs_pass,
        cs_directory=cs_directory) as cs:

        beacons = cs.get_beacons()

        print("[*] Getting beacon logs from teamserver...")
        beaconsresult = beacons

    ####################
    ## Process Logs

    # JSON field reference: type, beacon_id, user, command, result, timestamp

    if beaconsresult is None:
        print("[!] No logs yet. Did you just start the teamserver?")
        exit()

    links = []

    # Add Node Icons
    for beacon in beaconsresult:
        print(beacon)

        nodeIcon = u'\uf0e7'

        if beacon["pbid"] == "":
            nodeIcon = u'\uf0e7'

        else:
            nodeIcon = u'\uf0e7'

        beacon.update({"nodeIcon":nodeIcon})
        beacon.update({"build":str(beacon["build"])})

    # Create Links
    for beacon in beaconsresult:
        beacon_source = beacon["id"]
        beacon_target = ""
        beacon_type = ""

        if beacon["phint"] == "":
            beacon_type = "HTTP"
            beacon_target = "0" # teamserver
        elif beacon["phint"] == "445":
            beacon_type = "SMB"
            beacon_target = beacon["pbid"]
        else:
            beacon_type = "TCP"
            beacon_target = beacon["pbid"]


        # Add each beacon to list
        links.append({"source":beacon_source,"target":beacon_target,"type":beacon_type})

    # Add teamserver reference
    beaconsresult.append({
        'alive': 'true',
        'arch': '',
        'barch': '',
        'build': '0',
        'charset': '',
        'computer': '',
        'external': '',
        'host': 'teamserver',
        'id': '',
        'internal': '',
        'is64': '',
        'last': '',
        'lastf': '',
        'listener': '',
        'nodeIcon': '\uf0e7',
        'note': '',
        'os': 'Cobalt Strike',
        'pbid': '',
        'phint': '0',
        'pid': 'teamserver',
        'port': '',
        'process': 'teamserver',
        'session': '',
        'user': 'admin',
        'ver': 'teamserver',
        "nodeIcon":u'\uf233'
        })
    pprint(beaconsresult)
    pprint(links)

    output = json.dumps({"nodes":beacons,"links":links},ensure_ascii=False).encode('utf8')
    #print(output)

    with open('output/html/data/beacons.json', 'wb') as the_file:
        the_file.write(output)

if __name__ == "__main__":
    args = parseArguments()
    main(args)