class ArcticPhaseTests(unittest.TestCase):
# Create & Authenticate for subsequent tests
def setUp(self):
self.myatd = Atd('atd.localhost.localdomain')
error_control, data = self.myatd.connect('admin', 'password!')
self.assertEqual(1, error_control)
def test_config(self):
error_control, data = self.myatd.heartbeat()
self.assertEqual(1, error_control)
def test_config(self):
error_control, data = self.myatd.heartbeat()
self.assertEqual(1, error_control)
# Get the vmprofilelist
def test_md5(self):
error_control, data = self.myatd.get_vmprofiles()
self.assertEqual(1, error_control)
# Upload file to ATD Server
def test_upload(self):
error_control, data = self.myatd.upload_file('test/data/putty/putty_upx.exe', 24) # Make this a windows Profile
if error_control == 0:
self.assertFalse(True)
else:
# print '\nFile %s uploaded\n'%data['file']
self.assertIsInstance(data['jobId'], int)
self.assertIsInstance(data['taskId'], int)
self.assertIsInstance(data['md5'], unicode)
self.assertIsInstance(data['size'], int)
self.assertRegexpMatches(data['mimeType'], 'application')
def test_fetchreport(self):
error_control, data = self.myatd.get_report(8062)
if error_control == 1:
self.assertIsInstance(data['Summary']['Verdict']['Severity'], unicode)
self.assertIsInstance(data['Summary']['Verdict']['Description'], unicode)
else:
self.assertFalse(True)
if __name__ == '__main__':
unittest.main()
class SampleSubmit(CommonATD):
'''Class defining a file submission'''
def __init__(self, options):
# Create the ATD object and connect to it
self.rtnv = EXIT_FAILURE
self.options = options
# Get an authenticated connection to ATD
self.myatd = Atd(options.ip, options.skipssl)
self.error_control, self.data = self.myatd.connect(
self.options.user, self.options.password)
self.connection_check()
# Get the heartbeat value of the ATD Box
self.error_control, self.data = self.myatd.heartbeat()
self.heartbeat()
self.connection_check()
if self.options.verbosity:
print("T{1}SS file => {0}, id(options) {2}".format(
self.options.file_to_upload,
threading.currentThread().getName(), id(self.options)))
# Upload file to ATD Server
self.error_control, self.data = self.myatd.upload_file(
self.options.file_to_upload, self.options.analyzer_profile)
self.connection_check()
if self.options.verbosity > 2:
print(self.data)
job_id = self.data['jobId']
task_id = self.data['taskId']
if self.options.verbosity:
print('\nFile %s uploaded\n' % self.data['file'])
print('jobId: ', self.data['jobId'])
print('taskId: ', self.data['taskId'])
print('md5: ', self.data['md5'])
print('size: ', self.data['size'])
print('mimeType: ', self.data['mimeType'])
print('')
# Check status before requesting the report
stepwait = 5
while True:
error_control, self.data = self.myatd.check_status(task_id)
if error_control == 4 or error_control == 3:
if self.options.verbosity:
print('{0} - Waiting for {1} seconds'.format(
self.data, stepwait))
sys.stdout.flush()
else:
if self.options.quiet is not True:
print('.', end="")
sys.stdout.flush()
elif error_control == -1:
print(self.data)
self.myatd.disconnect()
sys.exit(-3)
else: # Analysis done
if self.options.verbosity:
print('\nAnalysis done')
else:
if self.options.quiet is not True:
print('.', end="")
sys.stdout.flush()
break
time.sleep(stepwait)
if stepwait < 30:
stepwait = stepwait + 5
# Getting Report information
if self.options.verbosity:
print('\nGetting report information...')
while True:
error_control, self.data = self.myatd.get_report(job_id)
if error_control == 0:
print('\n', self.data)
self.myatd.disconnect()
sys.exit(-4)
if error_control == 3:
print('\n', self.data)
self.myatd.disconnect()
sys.exit(0)
if error_control == 1:
try:
self.severity = self.data['Summary']['Verdict']['Severity']
self.md5 = self.data['Summary']['Subject']['md5']
if 'Description' in self.data['Summary']['Verdict']:
desc = self.data['Summary']['Verdict']['Description']
else:
desc = ""
except:
print('\n**BOMB parser**')
print(self.data)
self.myatd.disconnect()
sys.exit(-4)
else:
if self.options.quiet is not True:
print('({0}:{1}) = {2}: \"{3}\"'.format(
self.data['Summary']['Subject']['Name'],
self.data['Summary']['Subject']['md5'],
self.data['Summary']['Verdict']['Severity'], desc))
if self.options.verbosity:
print('\nFinal results...')
print(' Severity: %s' % self.severity)
print(' Description: %s' % desc)
if self.options.verbosity > 1:
print(self.data)
break
# error_control = 2
if self.options.verbosity:
print(' %s - Waiting for 30 seconds...' % self.data)
sys.stdout.flush()
time.sleep(30)
self.myatd.disconnect()
if self.options.quiet is not True:
print('')
self.rtnv = int(self.severity)
self.rtv_md5 = self.md5
# >>> DEBUG
# self.myatd.disconnect()
# self.rtnv = int(0)
# self.rtv_md5 = "HARDCODED"
# <<< END DEBUG
return
class ATDTests(unittest.TestCase):
# Create & Authenticate for subsequent tests
@vcr.use_cassette('fixtures/vcr_cassettes/connect.yaml')
def setUp(self):
self.myatd = Atd('atd.localhost.localdomain')
error_control, data = self.myatd.connect('admin', 'password.')
self.assertEqual(1, error_control)
# Test myatd properties
def test_connect(self):
self.assertIsInstance(self.myatd.session, unicode)
self.assertIsInstance(self.myatd.userId, unicode)
self.assertGreaterEqual(parse_version(self.myatd.matdver), parse_version(u'3.4.4.1.0'))
# Get the heartbeat value of the ATD Box
@vcr.use_cassette('fixtures/vcr_cassettes/heartbeat.yaml')
def test_heartbeat(self):
error_control, data = self.myatd.heartbeat()
self.assertEqual(1, error_control)
# Get the vmprofilelist
@vcr.use_cassette('fixtures/vcr_cassettes/vmprofiles.yaml')
def test_vmprofilelist(self):
error_control, data = self.myatd.get_vmprofiles()
self.assertEqual(1, error_control)
# Upload file to ATD Server
@vcr.use_cassette('fixtures/vcr_cassettes/upload_file.yaml')
def test_upload(self):
error_control, data = self.myatd.upload_file('test/data/putty/putty_upx.exe', 24) # Make this a windows Profile
if error_control == 0:
self.assertFalse(True)
else:
# print '\nFile %s uploaded\n'%data['file']
self.assertIsInstance(data['jobId'], int)
self.assertIsInstance(data['taskId'], int)
self.assertIsInstance(data['md5'], unicode)
self.assertIsInstance(data['size'], int)
self.assertRegexpMatches(data['mimeType'], 'application')
# Check status before requesting the report
@vcr.use_cassette('fixtures/vcr_cassettes/check_status.yaml')
def test_checkstatus_ok(self):
error_control, data = self.myatd.check_status(20079)
self.assertEqual(1, error_control)
# Check status before requesting the report
@vcr.use_cassette('fixtures/vcr_cassettes/check_status_notready.yaml')
def test_checkstatus_waiting(self):
error_control, data = self.myatd.check_status(20072)
self.assertIn(error_control, [1, 4])
@vcr.use_cassette('fixtures/vcr_cassettes/get_report.yaml')
def test_fetchreport(self):
error_control, data = self.myatd.get_report(8062)
if error_control == 1:
self.assertIsInstance(data['Summary']['Verdict']['Severity'], unicode)
self.assertIsInstance(data['Summary']['Verdict']['Description'], unicode)
else:
self.assertFalse(True)
@vcr.use_cassette('fixtures/vcr_cassettes/disconnect.yaml')
def test_disconnect(self):
error_control, data = self.myatd.disconnect()
# return(1,'Disconnection successful')
self.assertEqual(1, error_control)
if __name__ == '__main__':
unittest.main()