#!/usr/bin/env python
# This Source Code Form is subject to the terms of the Mozilla Public License,
# v. 2.0. If a copy of the MPL was not distributed with this file, You can
# obtain one at http://mozilla.org/MPL/2.0/.

import base64
import json
import logging
import os
import requests

import stoneridge


class StoneRidgePcap(object):
    def __init__(self, stop):
        self.stop = stop
        self.macaddr = stoneridge.get_config('machine', 'macaddr')
        self.netconfig = stoneridge.get_config('run', 'netconfig')
        self.host = stoneridge.get_config('tcpdump', self.netconfig)
        self.outdir = stoneridge.get_config('run', 'out')
        self.stdout = os.path.join(self.outdir, 'tcpdump.out')
        self.pcap = os.path.join(self.outdir, 'traffic.pcap')

        logging.debug('stop: %s' % (stop,))
        logging.debug('macaddr: %s' % (self.macaddr,))
        logging.debug('netconfig: %s' % (self.netconfig,))
        logging.debug('host: %s' % (self.host,))
        logging.debug('output directory: %s' % (self.outdir,))
        logging.debug('tcpdump output: %s' % (self.stdout,))
        logging.debug('packet capture: %s' % (self.pcap,))

    def start_pcap(self):
        url = 'http://%s/start/%s' % (self.host, self.macaddr)
        response = requests.post(url)
        res = json.loads(response.text)
        if res['status'] != 'ok':
            logging.error('Error starting pcap: %s' % (res['message'],))
        else:
            logging.debug('Started pcap')

    def stop_pcap(self):
        url = 'http://%s/stop/%s' % (self.host, self.macaddr)
        response = requests.post(url)
        res = json.loads(response.text)
        if res['status'] != 'ok':
            logging.error('Error stopping pcap: %s' % (res['message'],))
            return

        url = 'http://%s/retrieve/%s' % (self.host, self.macaddr)
        response = requests.post(url)
        res = json.loads(response.text)
        if res['status'] != 'ok':
            logging.error('Error retrieving pcap: %s' % (res['message'],))
            return

        stdout = base64.b64decode(res['data']['stdout'])
        pcap = base64.b64decode(res['data']['pcap'])

        with file(self.stdout, 'wb') as f:
            f.write(stdout)

        with file(self.pcap, 'wb') as f:
            f.write(pcap)

    def run(self):
        if self.stop:
            self.stop_pcap()
        else:
            self.start_pcap()


@stoneridge.main
def main():
    parser = stoneridge.TestRunArgumentParser()
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--start', dest='start', action='store_true')
    group.add_argument('--stop', dest='stop', action='store_true')
    args = parser.parse_args()

    pcap = StoneRidgePcap(args.stop)
    pcap.run()

@@ -1184,32 +1184,71 @@
 art_pcap(self):%0A
+        logging.debug('Starting pcap')%0A
         url = 'h
@@ -1324,32 +1324,208 @@
 uests.post(url)%0A
+        if response.status_code != 200:%0A            logging.error('Non-200 response starting pcap: %25s' %25%0A                          (response.status_code,))%0A            return%0A%0A
         res = js
@@ -1729,32 +1729,71 @@
 top_pcap(self):%0A
+        logging.debug('Stopping pcap')%0A
         url = 'h
@@ -1868,32 +1868,208 @@
 uests.post(url)%0A
+        if response.status_code != 200:%0A            logging.error('Non-200 response stopping pcap: %25s' %25%0A                          (response.status_code,))%0A            return%0A%0A
         res = js
@@ -2211,32 +2211,74 @@
         return%0A%0A
+        logging.debug('Retrieving pcap')%0A%0A
         url = 'h
@@ -2357,32 +2357,210 @@
 uests.post(url)%0A
+        if response.status_code != 200:%0A            logging.error('Non-200 response retrieving pcap: %25s' %25%0A                          (response.status_code,))%0A            return%0A%0A
         res = js
@@ -2882,65 +2882,203 @@
     
-f.write(stdout)%0A%0A        with file(self.pcap, 'wb') as f:
+logging.debug('Writing tcpdump stdout to %25s' %25 (self.stdout,))%0A            f.write(stdout)%0A%0A        with file(self.pcap, 'wb') as f:%0A            logging.debug('Writing pcap to %25s' %25 (self.pcap,))
 %0A   

Add more logging (and error handling) to client-side pcap