blob: dcc03cb64778114a5502049d670c2d1b57abc87a [file] [log] [blame]
koder aka kdanilov652cd802015-04-13 12:21:07 +03001import re
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +03002import os
koder aka kdanilovafd98742015-04-24 01:27:22 +03003import socket
koder aka kdanilove21d7472015-02-14 19:02:04 -08004import logging
koder aka kdanilov7acd6bd2015-02-12 14:28:30 -08005import threading
6import contextlib
koder aka kdanilov652cd802015-04-13 12:21:07 +03007import subprocess
koder aka kdanilov4643fd62015-02-10 16:20:13 -08008
koder aka kdanilove21d7472015-02-14 19:02:04 -08009
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +030010logger = logging.getLogger("wally")
koder aka kdanilov7acd6bd2015-02-12 14:28:30 -080011
12
koder aka kdanilove06762a2015-03-22 23:32:09 +020013def parse_creds(creds):
14 # parse user:passwd@host
15 user, passwd_host = creds.split(":", 1)
16
17 if '@' not in passwd_host:
18 passwd, host = passwd_host, None
19 else:
20 passwd, host = passwd_host.rsplit('@', 1)
21
22 return user, passwd, host
23
24
koder aka kdanilov2c473092015-03-29 17:12:13 +030025class TaksFinished(Exception):
26 pass
koder aka kdanilov4643fd62015-02-10 16:20:13 -080027
koder aka kdanilov2c473092015-03-29 17:12:13 +030028
koder aka kdanilove2de58c2015-04-24 22:59:36 +030029class StopTestError(RuntimeError):
30 def __init__(self, reason, orig_exc=None):
31 RuntimeError.__init__(self, reason)
32 self.orig_exc = orig_exc
33
34
koder aka kdanilov2c473092015-03-29 17:12:13 +030035class Barrier(object):
36 def __init__(self, count):
37 self.count = count
38 self.curr_count = 0
39 self.cond = threading.Condition()
40 self.exited = False
41
42 def wait(self, timeout=None):
43 with self.cond:
44 if self.exited:
45 raise TaksFinished()
46
47 self.curr_count += 1
48 if self.curr_count == self.count:
49 self.curr_count = 0
50 self.cond.notify_all()
koder aka kdanilov652cd802015-04-13 12:21:07 +030051 return True
koder aka kdanilov4643fd62015-02-10 16:20:13 -080052 else:
koder aka kdanilov2c473092015-03-29 17:12:13 +030053 self.cond.wait(timeout=timeout)
koder aka kdanilov652cd802015-04-13 12:21:07 +030054 return False
koder aka kdanilov4643fd62015-02-10 16:20:13 -080055
koder aka kdanilov2c473092015-03-29 17:12:13 +030056 def exit(self):
57 with self.cond:
58 self.exited = True
koder aka kdanilov7acd6bd2015-02-12 14:28:30 -080059
60
61@contextlib.contextmanager
62def log_error(action, types=(Exception,)):
63 if not action.startswith("!"):
koder aka kdanilove21d7472015-02-14 19:02:04 -080064 logger.debug("Starts : " + action)
koder aka kdanilov7acd6bd2015-02-12 14:28:30 -080065 else:
66 action = action[1:]
67
68 try:
69 yield
70 except Exception as exc:
71 if isinstance(exc, types) and not isinstance(exc, StopIteration):
koder aka kdanilovec1b9732015-04-23 20:43:29 +030072 templ = "Error during {0} stage: {1!s}"
73 logger.debug(templ.format(action, exc))
koder aka kdanilov7acd6bd2015-02-12 14:28:30 -080074 raise
75
76
koder aka kdanilov2c473092015-03-29 17:12:13 +030077SMAP = dict(k=1024, m=1024 ** 2, g=1024 ** 3, t=1024 ** 4)
koder aka kdanilov8ad6e812015-03-22 14:42:18 +020078
79
80def ssize_to_b(ssize):
81 try:
koder aka kdanilov2c473092015-03-29 17:12:13 +030082 ssize = ssize.lower()
koder aka kdanilov8ad6e812015-03-22 14:42:18 +020083
koder aka kdanilov2c473092015-03-29 17:12:13 +030084 if ssize.endswith("b"):
85 ssize = ssize[:-1]
86 if ssize[-1] in SMAP:
87 return int(ssize[:-1]) * SMAP[ssize[-1]]
koder aka kdanilov8ad6e812015-03-22 14:42:18 +020088 return int(ssize)
89 except (ValueError, TypeError, AttributeError):
koder aka kdanilov2e928022015-04-08 13:47:15 +030090 raise ValueError("Unknow size format {0!r}".format(ssize))
koder aka kdanilov652cd802015-04-13 12:21:07 +030091
92
93def get_ip_for_target(target_ip):
koder aka kdanilovafd98742015-04-24 01:27:22 +030094 if not re.match("[0-9]+\.[0-9]+\.[0-9]+\.[0-9]$", target_ip):
95 target_ip = socket.gethostbyname(target_ip)
96
97 if target_ip in ('localhost', '127.0.0.1', '127.0.1.1'):
98 return '127.0.0.1'
99
koder aka kdanilov652cd802015-04-13 12:21:07 +0300100 cmd = 'ip route get to'.split(" ") + [target_ip]
101 data = subprocess.Popen(cmd, stdout=subprocess.PIPE).stdout.read()
102
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300103 rr1 = r'{0} via [.0-9]+ dev (?P<dev>.*?) src (?P<ip>[.0-9]+)$'
104 rr1 = rr1.replace(" ", r'\s+')
105 rr1 = rr1.format(target_ip.replace('.', r'\.'))
106
107 rr2 = r'{0} dev (?P<dev>.*?) src (?P<ip>[.0-9]+)$'
108 rr2 = rr2.replace(" ", r'\s+')
109 rr2 = rr2.format(target_ip.replace('.', r'\.'))
koder aka kdanilov652cd802015-04-13 12:21:07 +0300110
111 data_line = data.split("\n")[0].strip()
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300112 res1 = re.match(rr1, data_line)
113 res2 = re.match(rr2, data_line)
koder aka kdanilov652cd802015-04-13 12:21:07 +0300114
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300115 if res1 is not None:
116 return res1.group('ip')
koder aka kdanilov652cd802015-04-13 12:21:07 +0300117
koder aka kdanilovcff7b2e2015-04-18 20:48:15 +0300118 if res2 is not None:
119 return res2.group('ip')
120
121 raise OSError("Can't define interface for {0}".format(target_ip))
122
123
124def open_for_append_or_create(fname):
125 if not os.path.exists(fname):
126 return open(fname, "w")
127
128 fd = open(fname, 'r+')
129 fd.seek(0, os.SEEK_END)
130 return fd
131
132
133def sec_to_str(seconds):
134 h = seconds // 3600
135 m = (seconds % 3600) // 60
136 s = seconds % 60
137 return "{0}:{1:02d}:{2:02d}".format(h, m, s)
koder aka kdanilov168f6092015-04-19 02:33:38 +0300138
139
140def yamable(data):
141 if isinstance(data, (tuple, list)):
142 return map(yamable, data)
143
144 if isinstance(data, unicode):
145 return str(data)
146
147 if isinstance(data, dict):
148 res = {}
149 for k, v in data.items():
150 res[yamable(k)] = yamable(v)
151 return res
152
153 return data