1#!/usr/bin/env python
2# Copyright 2012 Google Inc. All Rights Reserved.
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""Simulate network characteristics directly in Python.
17
18Allows running replay without dummynet.
19"""
20
21import logging
22import platformsettings
23import re
24import time
25
26
27TIMER = platformsettings.timer
28
29
30class ProxyShaperError(Exception):
31  """Module catch-all error."""
32  pass
33
34class BandwidthValueError(ProxyShaperError):
35  """Raised for unexpected dummynet-style bandwidth value."""
36  pass
37
38
39class RateLimitedFile(object):
40  """Wrap a file like object with rate limiting.
41
42  TODO(slamm): Simulate slow-start.
43      Each RateLimitedFile corresponds to one-direction of a
44      bidirectional socket. Slow-start can be added here (algorithm needed).
45      Will consider changing this class to take read and write files and
46      corresponding bit rates for each.
47  """
48  BYTES_PER_WRITE = 1460
49
50  def __init__(self, request_counter, f, bps):
51    """Initialize a RateLimiter.
52
53    Args:
54      request_counter: callable to see how many requests share the limit.
55      f: file-like object to wrap.
56      bps: an integer of bits per second.
57    """
58    self.request_counter = request_counter
59    self.original_file = f
60    self.bps = bps
61
62  def transfer_seconds(self, num_bytes):
63    """Seconds to read/write |num_bytes| with |self.bps|."""
64    return 8.0 * num_bytes / self.bps
65
66  def write(self, data):
67    num_bytes = len(data)
68    num_sent_bytes = 0
69    while num_sent_bytes < num_bytes:
70      num_write_bytes = min(self.BYTES_PER_WRITE, num_bytes - num_sent_bytes)
71      num_requests = self.request_counter()
72      wait = self.transfer_seconds(num_write_bytes) * num_requests
73      logging.debug('write sleep: %0.4fs (%d requests)', wait, num_requests)
74      time.sleep(wait)
75
76      self.original_file.write(
77          data[num_sent_bytes:num_sent_bytes + num_write_bytes])
78      num_sent_bytes += num_write_bytes
79
80  def _read(self, read_func, size):
81    start = TIMER()
82    data = read_func(size)
83    read_seconds = TIMER() - start
84    num_bytes = len(data)
85    num_requests = self.request_counter()
86    wait = self.transfer_seconds(num_bytes) * num_requests - read_seconds
87    if wait > 0:
88      logging.debug('read sleep: %0.4fs %d requests)', wait, num_requests)
89      time.sleep(wait)
90    return data
91
92  def readline(self, size=-1):
93    return self._read(self.original_file.readline, size)
94
95  def read(self, size=-1):
96    return self._read(self.original_file.read, size)
97
98  def __getattr__(self, name):
99    """Forward any non-overriden calls."""
100    return getattr(self.original_file, name)
101
102
103def GetBitsPerSecond(bandwidth):
104  """Return bits per second represented by dummynet bandwidth option.
105
106  See ipfw/dummynet.c:read_bandwidth for how it is really done.
107
108  Args:
109    bandwidth: a dummynet-style bandwidth specification (e.g. "10Kbit/s")
110  """
111  if bandwidth == '0':
112    return 0
113  bw_re = r'^(\d+)(?:([KM])?(bit|Byte)/s)?$'
114  match = re.match(bw_re, str(bandwidth))
115  if not match:
116    raise BandwidthValueError('Value, "%s", does not match regex: %s' % (
117        bandwidth, bw_re))
118  bw = int(match.group(1))
119  if match.group(2) == 'K':
120    bw *= 1000
121  if match.group(2) == 'M':
122    bw *= 1000000
123  if match.group(3) == 'Byte':
124    bw *= 8
125  return bw
126