1import os
2import re
3import csv
4import tempfile
5from datetime import datetime
6from collections import defaultdict
7from itertools import izip_longest
8
9from devlib.instrument import Instrument, MeasurementsCsv, CONTINUOUS
10from devlib.exception import TargetError, HostError
11from devlib.utils.android import ApkInfo
12
13
14THIS_DIR = os.path.dirname(__file__)
15
16NETSTAT_REGEX = re.compile(r'I/(?P<tag>netstats-\d+)\(\s*\d*\): (?P<ts>\d+) '
17                           r'"(?P<package>[^"]+)" TX: (?P<tx>\S+) RX: (?P<rx>\S+)')
18
19
20def extract_netstats(filepath, tag=None):
21    netstats = []
22    with open(filepath) as fh:
23        for line in fh:
24            match = NETSTAT_REGEX.search(line)
25            if not match:
26                continue
27            if tag and match.group('tag') != tag:
28                continue
29            netstats.append((match.group('tag'),
30                             match.group('ts'),
31                             match.group('package'),
32                             match.group('tx'),
33                             match.group('rx')))
34    return netstats
35
36
37def netstats_to_measurements(netstats):
38    measurements = defaultdict(list)
39    for row in netstats:
40        tag, ts, package, tx, rx = row  # pylint: disable=unused-variable
41        measurements[package + '_tx'].append(tx)
42        measurements[package + '_rx'].append(rx)
43    return measurements
44
45
46def write_measurements_csv(measurements, filepath):
47    headers = sorted(measurements.keys())
48    columns = [measurements[h] for h in headers]
49    with open(filepath, 'wb') as wfh:
50        writer = csv.writer(wfh)
51        writer.writerow(headers)
52        writer.writerows(izip_longest(*columns))
53
54
55class NetstatsInstrument(Instrument):
56
57    mode = CONTINUOUS
58
59    def __init__(self, target, apk=None, service='.TrafficMetricsService'):
60        """
61        Additional paramerter:
62
63        :apk: Path to the APK file that contains ``com.arm.devlab.netstats``
64              package. If not specified, it will be assumed that an APK with
65              name "netstats.apk" is located in the same directory as the
66              Python module for the instrument.
67        :service: Name of the service to be launched. This service must be
68                  present in the APK.
69
70        """
71        if target.os != 'android':
72            raise TargetError('netstats insturment only supports Android targets')
73        if apk is None:
74            apk = os.path.join(THIS_DIR, 'netstats.apk')
75        if not os.path.isfile(apk):
76            raise HostError('APK for netstats instrument does not exist ({})'.format(apk))
77        super(NetstatsInstrument, self).__init__(target)
78        self.apk = apk
79        self.package = ApkInfo(self.apk).package
80        self.service = service
81        self.tag = None
82        self.command = None
83        self.stop_command = 'am kill {}'.format(self.package)
84
85        for package in self.target.list_packages():
86            self.add_channel(package, 'tx')
87            self.add_channel(package, 'rx')
88
89    def setup(self, force=False, *args, **kwargs):
90        if self.target.package_is_installed(self.package):
91            if force:
92                self.logger.debug('Re-installing {} (forced)'.format(self.package))
93                self.target.uninstall_package(self.package)
94                self.target.install(self.apk)
95            else:
96                self.logger.debug('{} already present on target'.format(self.package))
97        else:
98            self.logger.debug('Deploying {} to target'.format(self.package))
99            self.target.install(self.apk)
100
101    def reset(self, sites=None, kinds=None, channels=None, period=None):  # pylint: disable=arguments-differ
102        super(NetstatsInstrument, self).reset(sites, kinds, channels)
103        period_arg, packages_arg = '', ''
104        self.tag = 'netstats-{}'.format(datetime.now().strftime('%Y%m%d%H%M%s'))
105        tag_arg = ' --es tag {}'.format(self.tag)
106        if sites:
107            packages_arg = ' --es packages {}'.format(','.join(sites))
108        if period:
109            period_arg = ' --ei period {}'.format(period)
110        self.command = 'am startservice{}{}{} {}/{}'.format(tag_arg,
111                                                            period_arg,
112                                                            packages_arg,
113                                                            self.package,
114                                                            self.service)
115        self.target.execute(self.stop_command)  # ensure the service is not running.
116
117    def start(self):
118        if self.command is None:
119            raise RuntimeError('reset() must be called before start()')
120        self.target.execute(self.command)
121
122    def stop(self):
123        self.target.execute(self.stop_command)
124
125    def get_data(self, outfile):
126        raw_log_file = tempfile.mktemp()
127        self.target.dump_logcat(raw_log_file)
128        data = extract_netstats(raw_log_file)
129        measurements = netstats_to_measurements(data)
130        write_measurements_csv(measurements, outfile)
131        os.remove(raw_log_file)
132        return MeasurementsCsv(outfile, self.active_channels)
133
134    def teardown(self):
135        self.target.uninstall_package(self.package)
136