def testPostResults(client, sample_Job, sample_Agent, sample_valid_JWT):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
agent["jobsRunning"].append(job)
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent.save()
# send job results to the api server
response = client.post("/agent/history",
headers={"Content-Type": "application/json",
"Agent-ID": sample_Agent["agentID"]},
data=json.dumps({"job": job.to_json()}))
assert response.status_code == 200
assert response.json["success"] == "successfully saved job response"
# get finished jobs for the sample agent from the api server
response = client.get("/agent/history",
headers={"Content-Type": "application/json",
"Authorization": "Bearer " + sample_valid_JWT},
data=json.dumps({"agentID": sample_Agent["agentID"]}))
assert response.status_code == 200
# make sure all job fields were included from the sample job
assert len(jobsHistory := response.json["jobs"]) == 1
finishedJob = json.loads(jobsHistory[0])
assert finishedJob["executor"] == job["executor"]
assert finishedJob["filename"] == job["filename"]
assert finishedJob["description"] == job["description"]
assert finishedJob["os"] == job["os"]
assert finishedJob["user"] == job["user"]
timeCreated = timestampToDatetime(job["timeCreated"])
timeDispatched = timestampToDatetime(job["timeDispatched"])
timeStarted = timestampToDatetime(job["timeStarted"])
timeEnded = timestampToDatetime(job["timeEnded"])
assert datetime.utcnow() >= timeCreated
assert timeDispatched >= timeCreated
assert timeStarted >= timeDispatched
assert timeEnded >= timeStarted
assert finishedJob["argv"] == job["argv"]
assert finishedJob["status"] == job["status"]
assert finishedJob["stdout"] == job["stdout"]
assert finishedJob["stderr"] == job["stderr"]
# post results again to verify that job was deleted from running queue
response = client.post("/agent/history",
headers={"Content-Type": "application/json",
"Agent-ID": sample_Agent["agentID"]},
data=json.dumps({"job": job.to_json()}))
assert response.status_code == 400
assert response.json["error"] == "no matching jobs were supposed to be running"
def markSent(self, agentID):
""" Remove first job from agent's cache and transfer from jobsQueue to jobsRunning in DB. """
while self.mutex:
sleep(1)
self.mutex = True
# job is available, return it from the cache
agent = Agent.objects(agentID=agentID).first()
jobsQueue = sorted(agent["jobsQueue"], key=lambda i: i["timeCreated"])
job = jobsQueue.pop(0)
job["timeDispatched"] = utcNowTimestamp()
agent["jobsRunning"].append(job)
agent["lastCheckin"] = utcNowTimestamp()
agent.save()
self.jobAssignments["agents"][agentID].pop(0)
self.mutex = False
def sample_Job():
job = Job(executor="psh",
filename="testfile",
description="Test job description. This job is not real.",
os="windows",
user="******",
timeCreated=utcNowTimestamp())
return job
def sample_Agent():
agent = Agent(hostname="testhost",
agentID=str(uuid4()),
os="windows",
lastCheckin=utcNowTimestamp(),
jobsQueue=[],
jobsRunning=[],
jobsHistory=[])
return agent
def testMissingFieldsPostResults(client, sample_Job, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
agent["jobsRunning"].append(job)
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent.save()
# send job results to the api server
response = client.post("/agent/history",
headers={"Content-Type": "application/json"},
data=json.dumps({}))
assert response.status_code == 400
assert response.json["error"] == "request is missing the following parameters: headers=['Agent-ID'], data=['job']"
def testGetResults(client, sample_Job, sample_Agent, sample_valid_JWT):
# prepare mongomock with relevant sample documents
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent = sample_Agent
agent["jobsHistory"].append(job)
agent.save()
# get finished jobs for sample agent from the api server
response = client.get("/agent/history",
headers={
"Content-Type": "application/json",
"Authorization": "Bearer " + sample_valid_JWT
},
data=json.dumps({"agentID":
sample_Agent["agentID"]}))
assert response.status_code == 200
# make sure all job fields were included from the sample job
assert len(jobsHistory := response.json["jobs"]) == 1
finishedJob = json.loads(jobsHistory[0])
assert finishedJob["executor"] == job["executor"]
assert finishedJob["filename"] == job["filename"]
assert finishedJob["description"] == job["description"]
assert finishedJob["os"] == job["os"]
assert finishedJob["user"] == job["user"]
timeCreated = timestampToDatetime(job["timeCreated"])
timeDispatched = timestampToDatetime(job["timeDispatched"])
timeStarted = timestampToDatetime(job["timeStarted"])
timeEnded = timestampToDatetime(job["timeEnded"])
assert datetime.utcnow() >= timeCreated
assert timeDispatched >= timeCreated
assert timeStarted >= timeDispatched
assert timeEnded >= timeStarted
assert finishedJob["argv"] == job["argv"]
assert finishedJob["status"] == job["status"]
assert finishedJob["stdout"] == job["stdout"]
assert finishedJob["stderr"] == job["stderr"]
def testMissingJobPostResults(client, sample_Job, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
# intentionally not adding job to agent's running queue
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent.save()
# send job results to the api server
response = client.post("/agent/history",
headers={"Content-Type": "application/json",
"Agent-ID": sample_Agent["agentID"]},
data=json.dumps({"job": job.to_json()}))
assert response.status_code == 400
assert response.json["error"] == "no matching jobs were supposed to be running"
def testUnknownAgentPostResults(client, sample_Job, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
agent["jobsRunning"].append(job)
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent.save()
# send job results to the api server
response = client.post("/agent/history",
headers={"Content-Type": "application/json",
"Agent-ID": "not_an_agent"},
data=json.dumps({"job": job.to_json()}))
assert response.status_code == 400
assert response.json["error"] == "agent ID not found"
def testUnknownAgentGetResults(client, sample_Job, sample_Agent,
sample_valid_JWT):
# prepare mongomock with relevant sample documents
job = sample_Job
job["timeDispatched"] = utcNowTimestamp()
job.argv = ["-o", "output.txt", "-i", "input.txt"]
job["timeStarted"] = utcNowTimestamp()
job["status"] = 0
job["stdout"] = "stdout"
job["stderr"] = "stderr"
job["timeEnded"] = utcNowTimestamp()
agent = sample_Agent
agent["jobsHistory"].append(job)
agent.save()
# get finished jobs for sample agent from the api server
response = client.get("/agent/history",
headers={
"Content-Type": "application/json",
"Authorization": "Bearer " + sample_valid_JWT
},
data=json.dumps({"agentID": "not_an_agent"}))
assert response.status_code == 400
assert response.json["error"] == "agent ID not found"
def testMissingFieldsGetExecutable(client, sample_Job, sample_JobFile,
sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job.filename = sample_JobFile
job.timeDispatched = utcNowTimestamp()
agent["jobsRunning"].append(job)
agent.save()
# get job's executable from the api server
response = client.get("/agent/execute",
headers={"Content-Type": "application/json"},
data=json.dumps({}))
assert response.status_code == 400
assert response.json[
"error"] == "request is missing the following parameters: headers=['Agent-ID'], data=['filename']"
def testNoJobGetExecutable(client, sample_Job, sample_JobFile, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job.filename = sample_JobFile
job.timeDispatched = utcNowTimestamp()
# intentionally not adding job to agent's jobsRunning list
agent.save()
# check in with api server
response = client.get("/agent/execute",
headers={
"Content-Type": "application/json",
"Agent-ID": sample_Agent.agentID
},
data=json.dumps({"filename": job.filename}))
assert response.status_code == 400
assert response.json["error"] == "no matching job available for download"
def testMissingFileGetExecutable(client, sample_Job, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
# intentionally not saving file to job library
job.timeDispatched = utcNowTimestamp()
agent["jobsRunning"].append(job)
agent.save()
# check in with api server
response = client.get("/agent/execute",
headers={
"Content-Type": "application/json",
"Agent-ID": sample_Agent.agentID
},
data=json.dumps({"filename": job.filename}))
assert response.status_code == 500
assert response.json[
"error"] == "job file missing -- please contact an administrator"
def testUnknownAgentGetExecutable(client, sample_Job, sample_JobFile,
sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job.filename = sample_JobFile
job.timeDispatched = utcNowTimestamp()
agent["jobsRunning"].append(job)
agent.save()
# check in with api server
response = client.get("/agent/execute",
headers={
"Content-Type": "application/json",
"Agent-ID": "not_an_agent"
},
data=json.dumps({"filename": job.filename}))
assert response.status_code == 400
assert response.json["error"] == "agent ID not found"
def testGetExecutable(client, sample_Job, sample_JobFile, sample_Agent):
# prepare mongomock with relevant sample documents
agent = sample_Agent
job = sample_Job
job.filename = sample_JobFile
job.timeDispatched = utcNowTimestamp()
agent["jobsRunning"].append(job)
agent.save()
# get job's executable from the api server
response = client.get("/agent/execute",
headers={
"Content-Type": "application/json",
"Agent-ID": sample_Agent.agentID
},
data=json.dumps({"filename": job.filename}))
assert response.status_code == 200
responseDisposition = response.headers["Content-Disposition"]
responseFilename = responseDisposition[responseDisposition.
index("filename=") + 9:]
assert responseFilename == job.filename
assert response.data == b'test content'
log.warning(f"<{request.remote_addr}> {missingParams}")
return {"error": missingParams}, 400
# check registration key
regKey = RegistrationKey.objects().first()
if not regKey:
log.error("no agent registration key found")
return {"error": "no registration key has been generated yet"}, 500
if regKey["regKey"] != request.json["registrationKey"]:
log.warning(f"<{request.remote_addr}> invalid registration key")
return {"error": "invalild registration key"}, 401
# TODO: make sure OS is valid
# create agent doc and add it to the db
newAgent = Agent(hostname=request.json["hostname"],
agentID=str(uuid4()),
os=request.json["os"],
lastCheckin=utcNowTimestamp())
newAgent.save()
# return agent ID
log.info(
f"<{request.remote_addr}> registered new agent {newAgent['agentID']}")
return {"agentID": newAgent["agentID"]}
@sock.route("/agent/checkin")
def agentCheckin(ws):
""" Agent checking in -- send file to be executed if a job is waiting """
log.debug(f"<{ws.sock.getpeername()[0]}> agent checking in")
# get agent ID from socket
while True:
data = ws.receive()
# validate that json data was sent