import virtualbox
import os
import time
import threading
import logging

logging.basicConfig(level=logging.DEBUG,
                    format='(%(threadName)-10s) %(message)s',)


class Sandbox(object):
    def __init__(self, name):
        vbox = virtualbox.VirtualBox()
        self.vm = vbox.find_machine(name)
        self.session = None
        self.gs = None
        self.malware_name = None

    def start(self):
        logging.info("Starting the VM")
        session = virtualbox.Session()
        power_up_process = self.vm.launch_vm_process(session, 'gui', '')
        power_up_process.wait_for_completion()
        session.unlock_machine()

    def power_off(self):
        logging.info("Shutdown the VM")
        self.gs.close()
        power_down_process = self.session.console.power_down()
        power_down_process.wait_for_completion()
        time.sleep(5)

    def restore_snapshot(self, point):
        logging.info("Restored the VM")
        session = self.vm.create_session()
        origin_point = session.machine.find_snapshot(point)
        restore_process = session.console.restore_snapshot(origin_point)
        restore_process.wait_for_completion()
        session.unlock_machine()

    def take_snapshot(self, point):
        logging.info("Take the snapshoot on %s" % (point))
        session = self.vm.create_session()
        localtime = time.asctime(time.localtime(time.time()))
        snapshot_process = self.session.console.take_snapshot(point, localtime)
        snapshot_process.wait_for_completion()
        session.unlock_machine()

    def delete_snapshot(self, point):
        logging.info("Delete the snapshoot %s" % (point))
        session = self.vm.create_session()
        origin_point = session.machine.find_snapshot(point)
        delete_process = session.console.delete_snapshot(origin_point.id_p)
        delete_process.wait_for_completion()
        session.unlock_machine()

    def capture_traffic(self, wait_sec, interface=0):
        if not os.path.isdir('PCAPLog'):
            os.mkdir('PCAPLog')
        logging.debug('Capture the network traffic')
        session = self.vm.create_session()
        adapter = session.machine.get_network_adapter(interface)
        adapter.trace_file = os.path.abspath('./PCAPLog/%s.pcap'
                                             % self.malware_name)
        adapter.trace_enabled = True
        time.sleep(wait_sec)
        adapter.trace_enabled = False
        session.unlock_machine()
        logging.debug('Exiting')


class Win32Sandbox(Sandbox):
    def __init__(self, name):
        super(Win32Sandbox, self).__init__(name)

    def __forward_timer(self, output, forward_hour):
        # output[0] == process; output[1] == stdout; output[2] == stderr
        current_time = output[1][5:-2]
        hour = int(current_time[0:2])
        if hour == 12:
            hour = 1
        else:
            hour = hour + forward_hour
        new_time = '{hh}:{mm}'.format(hh=str(hour), mm=current_time[3:5])
        logging.debug('Correct the time to: ' + new_time)
        return new_time

    def login(self, account, password):
        logging.info("Login to Windows XP")
        self.session = self.vm.create_session()
        self.gs = self.session.console.guest.create_session('john', 'john')
        time.sleep(10)
        try:
            self.gs.execute('C:\\Windows\\System32\\cmd.exe',
                            ['/C', 'tasklist'])
        except:
            logging.info("Waiting for logins to complete")
            time.sleep(15)

    def copy_malware_in(self, path):
        self.malware_name = path.split('/')[-1]
        logging.info("Copy malware %s into sandbox" % (self.malware_name))
        copy_process = self.gs.copy_to(path, 'C:\\%s.exe' % (self.malware_name))
        copy_process.wait_for_completion()

    def active_malware(self, name, timeout=5000):
        logging.debug('Activate the malware')
        try:
            process, stdout, stderr = self.gs.execute('C:\\%s.exe' % (name),
                                                      timeout_ms=timeout)
        except:
            f = open("./cannot_exec.out", "a")
            f.write(name + '\n')
            f.close()
            logging.info('Execute Error')
            logging.debug('Exiting')
            return

        output = self.gs.execute('C:\\Windows\\System32\\cmd.exe',
                                 ['/C', 'time', '/T'], timeout_ms=timeout)

        new_time = self.__forward_timer(output, 1)

        self.gs.execute('C:\\Windows\\System32\\cmd.exe',
                        ['/C', 'time', new_time], timeout_ms=timeout)

        output = self.gs.execute('C:\\Windows\\System32\\cmd.exe',
                                 ['/C', 'time', '/T'], timeout_ms=timeout)

        logging.debug('Exiting')

    def capture_traffic(self, wait_sec, interface=0):
        super(Win32Sandbox, self).capture_traffic(wait_sec, interface)

    def active_capture_parallel(self, time):
        capture_worker = threading.Thread(name='Capture Traffic',
                                          target=self.capture_traffic,
                                          args=(time, ))
        active_worker = threading.Thread(name='Active Malware',
                                         target=self.active_malware,
                                         args=(self.malware_name, ))
        capture_worker.start()
        active_worker.start()
        capture_worker.join()
        active_worker.join(5)


class LinuxSandbox(Sandbox):
    def __init__(self, name):
        super(LinuxSandbox, self).__init__(name)

    def bash_echo(self):
        # To be test on linux sandbox
        process, stdout, stderr = self.gs.execute('/bin/bash',
                                                  ['-c', 'echo hello world'])
        if stdout:
            logging.info(stdout)
        if stderr:
            logging.info(stderr)

    def capture_traffic(self, wait_sec, interface=0):
        super(LinuxSandbox, self).capture_traffic(wait_sec, interface)


if __name__ == "__main__":
    pass

@@ -1965,32 +1965,250 @@
 , interface=0):%0A
+        save_path = './PCAPLog/%7Bd%7D'.format(d=self.malware_name)%0A        save_file = './PCAPLog/%7Bd%7DLog%5C/%7Bf%7D.pcap'.format(d=self.malware_name,%0A                                                        f=self.malware_name)%0A
         if not o
@@ -2219,26 +2219,27 @@
 th.isdir
-('PCAPLog'
+s(save_path
 ):%0A     
@@ -2257,18 +2257,19 @@
 kdir
-('PCAPLog'
+s(save_path
 )%0A  
@@ -2476,91 +2476,16 @@
 ath(
-'./PCAPLog/%25s.pcap'%0A                                             %25 self.malware_nam
+save_fil
 e)%0A 

Add pcap log