Skip to content

Commit f257780

Browse files
maxDcbmaxDcb
committed
refactor: clean up teamserver grpc contract
1 parent ec7ef7b commit f257780

36 files changed

Lines changed: 677 additions & 449 deletions

‎C2Client/C2Client/ConsolePanel.py‎

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -590,43 +590,46 @@ def runCommand(self):
590590
self.commandEditor.setCmdHistory()
591591
instructions = commandLine.split()
592592
if instructions[0]==HelpInstruction:
593-
command = TeamServerApi_pb2.Command(
594-
beaconHash=self.beaconHash,
595-
listenerHash=self.listenerHash,
596-
cmd=commandLine,
597-
commandId=uuid.uuid4().hex,
593+
command = TeamServerApi_pb2.CommandHelpRequest(
594+
session=TeamServerApi_pb2.SessionSelector(
595+
beacon_hash=self.beaconHash,
596+
listener_hash=self.listenerHash,
597+
),
598+
command=commandLine,
598599
)
599-
response = self.grpcClient.getHelp(command)
600-
self.printInTerminal(response.cmd, "", "")
601-
self.printInTerminal("", response.cmd, response.response.decode('utf-8', 'replace'))
600+
response = self.grpcClient.getCommandHelp(command)
601+
self.printInTerminal(response.command, "", "")
602+
self.printInTerminal("", response.command, response.help)
602603

603604
else:
604605
self.printInTerminal(commandLine, "", "")
605606
command_id = uuid.uuid4().hex
606-
command = TeamServerApi_pb2.Command(
607-
beaconHash=self.beaconHash,
608-
listenerHash=self.listenerHash,
609-
cmd=commandLine,
610-
commandId=command_id,
607+
command = TeamServerApi_pb2.SessionCommandRequest(
608+
session=TeamServerApi_pb2.SessionSelector(
609+
beacon_hash=self.beaconHash,
610+
listener_hash=self.listenerHash,
611+
),
612+
command=commandLine,
613+
command_id=command_id,
611614
)
612-
result = self.grpcClient.sendCmdToSession(command)
613-
command_id = getattr(result, "commandId", command_id) or command_id
615+
result = self.grpcClient.sendSessionCommand(command)
616+
command_id = getattr(result, "command_id", command_id) or command_id
614617
context = "Host " + self.hostname + " - Username " + self.username
615618
self.consoleScriptSignal.emit("send", self.beaconHash, self.listenerHash, context, commandLine, "", command_id)
616619
if result.message:
617-
self.printInTerminal("", commandLine, result.message.decode('utf-8', 'replace'))
620+
self.printInTerminal("", commandLine, result.message)
618621

619622
self.setCursorEditorAtEnd()
620623

621624
def displayResponse(self):
622-
session = TeamServerApi_pb2.Session(beaconHash=self.beaconHash, listenerHash=self.listenerHash)
623-
responses = self.grpcClient.getResponseFromSession(session)
625+
session = TeamServerApi_pb2.SessionSelector(beacon_hash=self.beaconHash, listener_hash=self.listenerHash)
626+
responses = self.grpcClient.streamSessionCommandResults(session)
624627
for response in responses:
625628
context = "Host " + self.hostname + " - Username " + self.username
626-
command_id = getattr(response, "commandId", "")
627-
listener_hash = getattr(response, "listenerHash", "") or self.listenerHash
628-
command_text = response.cmd or response.instruction
629-
decoded_response = response.response.decode('utf-8', 'replace')
629+
command_id = getattr(response, "command_id", "")
630+
listener_hash = response.session.listener_hash or self.listenerHash
631+
command_text = response.command or response.instruction
632+
decoded_response = response.output.decode('utf-8', 'replace')
630633
self.consoleScriptSignal.emit("receive", self.beaconHash, listener_hash, context, command_text, decoded_response, command_id)
631634
self.setCursorEditorAtEnd()
632635
# check the response for mimikatz and not the cmd line ???

‎C2Client/C2Client/GraphPanel.py‎

Lines changed: 17 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -224,7 +224,7 @@ def updateGraph(self):
224224
#
225225
# Update beacons
226226
#
227-
responses = self.grpcClient.getSessions()
227+
responses = self.grpcClient.listSessions()
228228
sessions = list()
229229
for response in responses:
230230
sessions.append(response)
@@ -233,7 +233,7 @@ def updateGraph(self):
233233
for ix, nodeItem in enumerate(self.listNodeItem):
234234
runing=False
235235
for session in sessions:
236-
if session.beaconHash == nodeItem.beaconHash:
236+
if session.beacon_hash == nodeItem.beaconHash:
237237
runing=True
238238
if not runing and self.listNodeItem[ix].type == BeaconNodeItemType:
239239
for ix2, connector in enumerate(self.listConnector):
@@ -249,20 +249,20 @@ def updateGraph(self):
249249
for session in sessions:
250250
inStore=False
251251
for ix, nodeItem in enumerate(self.listNodeItem):
252-
if session.beaconHash == nodeItem.beaconHash:
252+
if session.beacon_hash == nodeItem.beaconHash:
253253
inStore=True
254254
if not inStore:
255-
item = NodeItem(BeaconNodeItemType, session.beaconHash, session.os, session.privilege, session.hostname)
256-
item.connectedListenerHash = session.listenerHash
255+
item = NodeItem(BeaconNodeItemType, session.beacon_hash, session.os, session.privilege, session.hostname)
256+
item.connectedListenerHash = session.listener_hash
257257
item.signaller.signal.connect(self.updateConnectors)
258258
self.scene.addItem(item)
259259
self.listNodeItem.append(item)
260-
print("[+] add beacon", session.beaconHash)
260+
print("[+] add beacon", session.beacon_hash)
261261

262262
#
263263
# Update listener
264264
#
265-
responses= self.grpcClient.getListeners()
265+
responses= self.grpcClient.listListeners()
266266
listeners = list()
267267
for listener in responses:
268268
listeners.append(listener)
@@ -271,7 +271,7 @@ def updateGraph(self):
271271
for ix, nodeItem in enumerate(self.listNodeItem):
272272
runing=False
273273
for listener in listeners:
274-
if nodeItem.isResponsableForListener(listener.listenerHash):
274+
if nodeItem.isResponsableForListener(listener.listener_hash):
275275
runing=True
276276
if not runing:
277277
# primary listener
@@ -287,33 +287,33 @@ def updateGraph(self):
287287

288288
# beacon listener
289289
elif self.listNodeItem[ix].type == BeaconNodeItemType:
290-
if listener.listenerHash in self.listNodeItem[ix].listenerHash:
290+
if listener.listener_hash in self.listNodeItem[ix].listenerHash:
291291
for ix2, connector in enumerate(self.listConnector):
292292
if self.listNodeItem[ix2].listenerHash in connector.listener.listenerHash:
293293
print("[-] delete connector")
294294
self.scene.removeItem(self.listConnector[ix2])
295295
del self.listConnector[ix2]
296296
print("[-] delete secondary listener", nodeItem.listenerHash)
297-
self.listNodeItem[ix].listenerHash.remove(listener.listenerHash)
297+
self.listNodeItem[ix].listenerHash.remove(listener.listener_hash)
298298

299299
# add listener
300300
for listener in listeners:
301301
inStore=False
302302
for ix, nodeItem in enumerate(self.listNodeItem):
303-
if nodeItem.isResponsableForListener(listener.listenerHash):
303+
if nodeItem.isResponsableForListener(listener.listener_hash):
304304
inStore=True
305305
if not inStore:
306-
if not listener.beaconHash:
307-
item = NodeItem(ListenerNodeItemType, listener.listenerHash)
306+
if not listener.beacon_hash:
307+
item = NodeItem(ListenerNodeItemType, listener.listener_hash)
308308
item.signaller.signal.connect(self.updateConnectors)
309309
self.scene.addItem(item)
310310
self.listNodeItem.append(item)
311-
print("[+] add primary listener", listener.listenerHash)
311+
print("[+] add primary listener", listener.listener_hash)
312312
else:
313313
for nodeItem2 in self.listNodeItem:
314-
if nodeItem2.beaconHash == listener.beaconHash:
315-
nodeItem2.listenerHash.append(listener.listenerHash)
316-
print("[+] add secondary listener", listener.listenerHash)
314+
if nodeItem2.beaconHash == listener.beacon_hash:
315+
nodeItem2.listenerHash.append(listener.listener_hash)
316+
print("[+] add secondary listener", listener.listener_hash)
317317

318318
#
319319
# Update connectors
@@ -361,4 +361,3 @@ def run(self):
361361

362362
def quit(self):
363363
self.exit=True
364-

‎C2Client/C2Client/ListenerPanel.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ def __init__(self, parent, grpcClient):
101101
self.getListenerWorker = GetListenerWorker()
102102
self.getListenerWorker.moveToThread(self.thread)
103103
self.thread.started.connect(self.getListenerWorker.run)
104-
self.getListenerWorker.checkin.connect(self.getListeners)
104+
self.getListenerWorker.checkin.connect(self.listListeners)
105105
self.thread.start()
106106

107107
self.setLayout(self.layout)
@@ -171,14 +171,14 @@ def addListener(self, message):
171171

172172
# send message for stoping a listener
173173
def stopListener(self, listenerHash):
174-
listener = TeamServerApi_pb2.Listener(
175-
listenerHash=listenerHash)
174+
listener = TeamServerApi_pb2.ListenerSelector(
175+
listener_hash=listenerHash)
176176
self.grpcClient.stopListener(listener)
177177

178178

179179
# query the server to get the list of listeners
180-
def getListeners(self):
181-
responses = self.grpcClient.getListeners()
180+
def listListeners(self):
181+
responses = self.grpcClient.listListeners()
182182

183183
listeners = list()
184184
for response in responses:
@@ -188,7 +188,7 @@ def getListeners(self):
188188
for ix, listenerStore in enumerate(self.listListenerObject):
189189
runing=False
190190
for listener in listeners:
191-
if listener.listenerHash == listenerStore.listenerHash:
191+
if listener.listener_hash == listenerStore.listenerHash:
192192
runing=True
193193
# delete
194194
if not runing:
@@ -199,23 +199,23 @@ def getListeners(self):
199199
# if listener is already on our list
200200
for ix, listenerStore in enumerate(self.listListenerObject):
201201
# maj
202-
if listener.listenerHash == listenerStore.listenerHash:
202+
if listener.listener_hash == listenerStore.listenerHash:
203203
inStore=True
204-
listenerStore.nbSession=listener.numberOfSession
204+
listenerStore.nbSession=listener.session_count
205205
# add
206206
# if listener is not yet already on our list
207207
if not inStore:
208208

209209
self.listenerScriptSignal.emit("start", "", "", "")
210210

211211
if listener.type == GithubType:
212-
self.listListenerObject.append(Listener(self.idListener, listener.listenerHash, listener.type, listener.project, listener.token[0:10], listener.numberOfSession))
212+
self.listListenerObject.append(Listener(self.idListener, listener.listener_hash, listener.type, listener.project, listener.token[0:10], listener.session_count))
213213
elif listener.type == DnsType:
214-
self.listListenerObject.append(Listener(self.idListener, listener.listenerHash, listener.type, listener.domain, listener.port, listener.numberOfSession))
214+
self.listListenerObject.append(Listener(self.idListener, listener.listener_hash, listener.type, listener.domain, listener.port, listener.session_count))
215215
elif listener.type == SmbType:
216-
self.listListenerObject.append(Listener(self.idListener, listener.listenerHash, listener.type, listener.ip, listener.domain, listener.numberOfSession))
216+
self.listListenerObject.append(Listener(self.idListener, listener.listener_hash, listener.type, listener.ip, listener.domain, listener.session_count))
217217
else:
218-
self.listListenerObject.append(Listener(self.idListener, listener.listenerHash, listener.type, listener.ip, listener.port, listener.numberOfSession))
218+
self.listListenerObject.append(Listener(self.idListener, listener.listener_hash, listener.type, listener.ip, listener.port, listener.session_count))
219219
self.idListener = self.idListener+1
220220

221221
self.printListeners()

‎C2Client/C2Client/Scripts/checkSandbox.py‎

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,14 @@ def OnSessionStart(grpcClient, beaconHash, listenerHash, hostname, username, arc
99
output += "checkSandbox:\nSandbox detected ending beacon\n";
1010

1111
commandLine = "end"
12-
command = TeamServerApi_pb2.Command(beaconHash=beaconHash, listenerHash=listenerHash, cmd=commandLine, commandId=uuid.uuid4().hex)
13-
result = grpcClient.sendCmdToSession(command)
12+
command = TeamServerApi_pb2.SessionCommandRequest(
13+
session=TeamServerApi_pb2.SessionSelector(
14+
beacon_hash=beaconHash,
15+
listener_hash=listenerHash,
16+
),
17+
command=commandLine,
18+
command_id=uuid.uuid4().hex,
19+
)
20+
result = grpcClient.sendSessionCommand(command)
1421

1522
return output

‎C2Client/C2Client/Scripts/listDirectory.py‎

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,18 @@ def OnSessionStart(grpcClient, beaconHash, listenerHash, hostname, username, arc
88
output += "load ListDirectory\n";
99

1010
commandLine = "loadModule ListDirectory"
11-
command = TeamServerApi_pb2.Command(beaconHash=beaconHash, listenerHash=listenerHash, cmd=commandLine, commandId=uuid.uuid4().hex)
12-
result = grpcClient.sendCmdToSession(command)
11+
command = TeamServerApi_pb2.SessionCommandRequest(
12+
session=TeamServerApi_pb2.SessionSelector(
13+
beacon_hash=beaconHash,
14+
listener_hash=listenerHash,
15+
),
16+
command=commandLine,
17+
command_id=uuid.uuid4().hex,
18+
)
19+
result = grpcClient.sendSessionCommand(command)
1320

1421
# commandLine = "ls"
15-
# command = TeamServerApi_pb2.Command(beaconHash=beaconHash, listenerHash=listenerHash, cmd=commandLine, commandId=uuid.uuid4().hex)
16-
# result = grpcClient.sendCmdToSession(command)
22+
# command = TeamServerApi_pb2.SessionCommandRequest(session=TeamServerApi_pb2.SessionSelector(beacon_hash=beaconHash, listener_hash=listenerHash), command=commandLine, command_id=uuid.uuid4().hex)
23+
# result = grpcClient.sendSessionCommand(command)
1724

1825
return output

‎C2Client/C2Client/SessionPanel.py‎

Lines changed: 23 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ def __init__(self, parent, grpcClient):
8181
self.getSessionsWorker = GetSessionsWorker()
8282
self.getSessionsWorker.moveToThread(self.thread)
8383
self.thread.started.connect(self.getSessionsWorker.run)
84-
self.getSessionsWorker.checkin.connect(self.getSessions)
84+
self.getSessionsWorker.checkin.connect(self.listSessions)
8585
self.thread.start()
8686

8787
self.setLayout(self.layout)
@@ -138,14 +138,14 @@ def actionClicked(self, action):
138138

139139

140140
def stopSession(self, beaconHash, listenerHash):
141-
session = TeamServerApi_pb2.Session(
142-
beaconHash=beaconHash, listenerHash=listenerHash)
141+
session = TeamServerApi_pb2.SessionSelector(
142+
beacon_hash=beaconHash, listener_hash=listenerHash)
143143
self.grpcClient.stopSession(session)
144-
self.getSessions()
144+
self.listSessions()
145145

146146

147-
def getSessions(self):
148-
responses = self.grpcClient.getSessions()
147+
def listSessions(self):
148+
responses = self.grpcClient.listSessions()
149149

150150
sessions = list()
151151
for response in responses:
@@ -155,7 +155,7 @@ def getSessions(self):
155155
for ix, item in enumerate(self.listSessionObject):
156156
runing=False
157157
for session in sessions:
158-
if session.beaconHash == item.beaconHash:
158+
if session.beacon_hash == item.beaconHash:
159159
runing=True
160160
# set idl
161161
if not runing:
@@ -165,11 +165,11 @@ def getSessions(self):
165165
inStore=False
166166
for sessionStore in self.listSessionObject:
167167
#maj
168-
if session.listenerHash == sessionStore.listenerHash and session.beaconHash == sessionStore.beaconHash:
169-
self.sessionScriptSignal.emit("update", session.beaconHash, session.listenerHash, session.hostname, session.username, session.arch, session.privilege, session.os, session.lastProofOfLife, session.killed)
168+
if session.listener_hash == sessionStore.listenerHash and session.beacon_hash == sessionStore.beaconHash:
169+
self.sessionScriptSignal.emit("update", session.beacon_hash, session.listener_hash, session.hostname, session.username, session.arch, session.privilege, session.os, session.last_proof_of_life, session.killed)
170170
inStore=True
171-
sessionStore.lastProofOfLife=session.lastProofOfLife
172-
sessionStore.listenerHash=session.listenerHash
171+
sessionStore.lastProofOfLife=session.last_proof_of_life
172+
sessionStore.listenerHash=session.listener_hash
173173
if session.hostname:
174174
sessionStore.hostname=session.hostname
175175
if session.username:
@@ -180,29 +180,29 @@ def getSessions(self):
180180
sessionStore.privilege=session.privilege
181181
if session.os:
182182
sessionStore.os=session.os
183-
if session.lastProofOfLife:
184-
sessionStore.lastProofOfLife=session.lastProofOfLife
183+
if session.last_proof_of_life:
184+
sessionStore.lastProofOfLife=session.last_proof_of_life
185185
if session.killed:
186186
sessionStore.killed=session.killed
187-
if session.internalIps:
188-
sessionStore.internalIps=session.internalIps
189-
if session.processId:
190-
sessionStore.processId=session.processId
191-
if session.additionalInformation:
192-
sessionStore.additionalInformation=session.additionalInformation
187+
if session.internal_ips:
188+
sessionStore.internalIps=session.internal_ips
189+
if session.process_id:
190+
sessionStore.processId=session.process_id
191+
if session.additional_information:
192+
sessionStore.additionalInformation=session.additional_information
193193
# add
194194
if not inStore:
195-
self.sessionScriptSignal.emit("start", session.beaconHash, session.listenerHash, session.hostname, session.username, session.arch, session.privilege, session.os, session.lastProofOfLife, session.killed)
195+
self.sessionScriptSignal.emit("start", session.beacon_hash, session.listener_hash, session.hostname, session.username, session.arch, session.privilege, session.os, session.last_proof_of_life, session.killed)
196196

197197
# print(session)
198198

199199
self.listSessionObject.append(
200200
Session(
201201
self.idSession,
202-
session.listenerHash, session.beaconHash,
202+
session.listener_hash, session.beacon_hash,
203203
session.hostname, session.username, session.arch,
204-
session.privilege, session.os, session.lastProofOfLife,
205-
session.killed, session.internalIps, session.processId, session.additionalInformation
204+
session.privilege, session.os, session.last_proof_of_life,
205+
session.killed, session.internal_ips, session.process_id, session.additional_information
206206
)
207207
)
208208
self.idSession = self.idSession+1

0 commit comments

Comments
 (0)