1"""
2Support for LVS (Linux Virtual Server)
3"""
4
5import salt.utils.decorators as decorators
6import salt.utils.path
7from salt.exceptions import SaltException
8
9__func_alias__ = {"list_": "list"}
10
11
12# Cache the output of running which('ipvsadm')
13@decorators.memoize
14def __detect_os():
15    return salt.utils.path.which("ipvsadm")
16
17
18def __virtual__():
19    """
20    Only load if ipvsadm command exists on the system.
21    """
22    if not __detect_os():
23        return (
24            False,
25            "The lvs execution module cannot be loaded: the ipvsadm binary is not in"
26            " the path.",
27        )
28
29    return "lvs"
30
31
32def _build_cmd(**kwargs):
33    """
34
35    Build a well-formatted ipvsadm command based on kwargs.
36    """
37    cmd = ""
38
39    if "service_address" in kwargs:
40        if kwargs["service_address"]:
41            if "protocol" in kwargs:
42                if kwargs["protocol"] == "tcp":
43                    cmd += " -t {}".format(kwargs["service_address"])
44                elif kwargs["protocol"] == "udp":
45                    cmd += " -u {}".format(kwargs["service_address"])
46                elif kwargs["protocol"] == "fwmark":
47                    cmd += " -f {}".format(kwargs["service_address"])
48                else:
49                    raise SaltException(
50                        "Error: Only support tcp, udp and fwmark service protocol"
51                    )
52                del kwargs["protocol"]
53            else:
54                raise SaltException("Error: protocol should specified")
55            if "scheduler" in kwargs:
56                if kwargs["scheduler"]:
57                    cmd += " -s {}".format(kwargs["scheduler"])
58                    del kwargs["scheduler"]
59        else:
60            raise SaltException("Error: service_address should specified")
61        del kwargs["service_address"]
62
63    if "server_address" in kwargs:
64        if kwargs["server_address"]:
65            cmd += " -r {}".format(kwargs["server_address"])
66            if "packet_forward_method" in kwargs and kwargs["packet_forward_method"]:
67                if kwargs["packet_forward_method"] == "dr":
68                    cmd += " -g"
69                elif kwargs["packet_forward_method"] == "tunnel":
70                    cmd += " -i"
71                elif kwargs["packet_forward_method"] == "nat":
72                    cmd += " -m"
73                else:
74                    raise SaltException("Error: only support dr, tunnel and nat")
75                del kwargs["packet_forward_method"]
76            if "weight" in kwargs and kwargs["weight"]:
77                cmd += " -w {}".format(kwargs["weight"])
78                del kwargs["weight"]
79        else:
80            raise SaltException("Error: server_address should specified")
81        del kwargs["server_address"]
82
83    return cmd
84
85
86def add_service(protocol=None, service_address=None, scheduler="wlc"):
87    """
88    Add a virtual service.
89
90    protocol
91        The service protocol(only support tcp, udp and fwmark service).
92
93    service_address
94        The LVS service address.
95
96    scheduler
97        Algorithm for allocating TCP connections and UDP datagrams to real servers.
98
99    CLI Example:
100
101    .. code-block:: bash
102
103        salt '*' lvs.add_service tcp 1.1.1.1:80 rr
104    """
105
106    cmd = "{} -A {}".format(
107        __detect_os(),
108        _build_cmd(
109            protocol=protocol, service_address=service_address, scheduler=scheduler
110        ),
111    )
112    out = __salt__["cmd.run_all"](cmd, python_shell=False)
113
114    # A non-zero return code means fail
115    if out["retcode"]:
116        ret = out["stderr"].strip()
117    else:
118        ret = True
119    return ret
120
121
122def edit_service(protocol=None, service_address=None, scheduler=None):
123    """
124    Edit the virtual service.
125
126    protocol
127        The service protocol(only support tcp, udp and fwmark service).
128
129    service_address
130        The LVS service address.
131
132    scheduler
133        Algorithm for allocating TCP connections and UDP datagrams to real servers.
134
135    CLI Example:
136
137    .. code-block:: bash
138
139        salt '*' lvs.edit_service tcp 1.1.1.1:80 rr
140    """
141
142    cmd = "{} -E {}".format(
143        __detect_os(),
144        _build_cmd(
145            protocol=protocol, service_address=service_address, scheduler=scheduler
146        ),
147    )
148    out = __salt__["cmd.run_all"](cmd, python_shell=False)
149
150    # A non-zero return code means fail
151    if out["retcode"]:
152        ret = out["stderr"].strip()
153    else:
154        ret = True
155    return ret
156
157
158def delete_service(protocol=None, service_address=None):
159    """
160
161    Delete the virtual service.
162
163    protocol
164        The service protocol(only support tcp, udp and fwmark service).
165
166    service_address
167        The LVS service address.
168
169    CLI Example:
170
171    .. code-block:: bash
172
173        salt '*' lvs.delete_service tcp 1.1.1.1:80
174    """
175
176    cmd = "{} -D {}".format(
177        __detect_os(), _build_cmd(protocol=protocol, service_address=service_address)
178    )
179    out = __salt__["cmd.run_all"](cmd, python_shell=False)
180
181    # A non-zero return code means fail
182    if out["retcode"]:
183        ret = out["stderr"].strip()
184    else:
185        ret = True
186    return ret
187
188
189def add_server(
190    protocol=None,
191    service_address=None,
192    server_address=None,
193    packet_forward_method="dr",
194    weight=1,
195    **kwargs
196):
197    """
198
199    Add a real server to a virtual service.
200
201    protocol
202        The service protocol(only support ``tcp``, ``udp`` and ``fwmark`` service).
203
204    service_address
205        The LVS service address.
206
207    server_address
208        The real server address.
209
210    packet_forward_method
211        The LVS packet forwarding method(``dr`` for direct routing, ``tunnel`` for tunneling, ``nat`` for network access translation).
212
213    weight
214        The capacity  of a server relative to the others in the pool.
215
216    CLI Example:
217
218    .. code-block:: bash
219
220        salt '*' lvs.add_server tcp 1.1.1.1:80 192.168.0.11:8080 nat 1
221    """
222
223    cmd = "{} -a {}".format(
224        __detect_os(),
225        _build_cmd(
226            protocol=protocol,
227            service_address=service_address,
228            server_address=server_address,
229            packet_forward_method=packet_forward_method,
230            weight=weight,
231            **kwargs
232        ),
233    )
234    out = __salt__["cmd.run_all"](cmd, python_shell=False)
235
236    # A non-zero return code means fail
237    if out["retcode"]:
238        ret = out["stderr"].strip()
239    else:
240        ret = True
241    return ret
242
243
244def edit_server(
245    protocol=None,
246    service_address=None,
247    server_address=None,
248    packet_forward_method=None,
249    weight=None,
250    **kwargs
251):
252    """
253
254    Edit a real server to a virtual service.
255
256    protocol
257        The service protocol(only support ``tcp``, ``udp`` and ``fwmark`` service).
258
259    service_address
260        The LVS service address.
261
262    server_address
263        The real server address.
264
265    packet_forward_method
266        The LVS packet forwarding method(``dr`` for direct routing, ``tunnel`` for tunneling, ``nat`` for network access translation).
267
268    weight
269        The capacity  of a server relative to the others in the pool.
270
271    CLI Example:
272
273    .. code-block:: bash
274
275        salt '*' lvs.edit_server tcp 1.1.1.1:80 192.168.0.11:8080 nat 1
276    """
277
278    cmd = "{} -e {}".format(
279        __detect_os(),
280        _build_cmd(
281            protocol=protocol,
282            service_address=service_address,
283            server_address=server_address,
284            packet_forward_method=packet_forward_method,
285            weight=weight,
286            **kwargs
287        ),
288    )
289    out = __salt__["cmd.run_all"](cmd, python_shell=False)
290
291    # A non-zero return code means fail
292    if out["retcode"]:
293        ret = out["stderr"].strip()
294    else:
295        ret = True
296    return ret
297
298
299def delete_server(protocol=None, service_address=None, server_address=None):
300    """
301
302    Delete the realserver from the virtual service.
303
304    protocol
305        The service protocol(only support ``tcp``, ``udp`` and ``fwmark`` service).
306
307    service_address
308        The LVS service address.
309
310    server_address
311        The real server address.
312
313    CLI Example:
314
315    .. code-block:: bash
316
317        salt '*' lvs.delete_server tcp 1.1.1.1:80 192.168.0.11:8080
318    """
319
320    cmd = "{} -d {}".format(
321        __detect_os(),
322        _build_cmd(
323            protocol=protocol,
324            service_address=service_address,
325            server_address=server_address,
326        ),
327    )
328    out = __salt__["cmd.run_all"](cmd, python_shell=False)
329
330    # A non-zero return code means fail
331    if out["retcode"]:
332        ret = out["stderr"].strip()
333    else:
334        ret = True
335    return ret
336
337
338def clear():
339    """
340
341    Clear the virtual server table
342
343    CLI Example:
344
345    .. code-block:: bash
346
347        salt '*' lvs.clear
348    """
349
350    cmd = "{} -C".format(__detect_os())
351
352    out = __salt__["cmd.run_all"](cmd, python_shell=False)
353
354    # A non-zero return code means fail
355    if out["retcode"]:
356        ret = out["stderr"].strip()
357    else:
358        ret = True
359    return ret
360
361
362def get_rules():
363    """
364
365    Get the virtual server rules
366
367    CLI Example:
368
369    .. code-block:: bash
370
371        salt '*' lvs.get_rules
372    """
373
374    cmd = "{} -S -n".format(__detect_os())
375
376    ret = __salt__["cmd.run"](cmd, python_shell=False)
377    return ret
378
379
380def list_(protocol=None, service_address=None):
381    """
382
383    List the virtual server table if service_address is not specified. If a service_address is selected, list this service only.
384
385    CLI Example:
386
387    .. code-block:: bash
388
389        salt '*' lvs.list
390    """
391
392    if service_address:
393        cmd = "{} -L {} -n".format(
394            __detect_os(),
395            _build_cmd(protocol=protocol, service_address=service_address),
396        )
397    else:
398        cmd = "{} -L -n".format(__detect_os())
399    out = __salt__["cmd.run_all"](cmd, python_shell=False)
400
401    # A non-zero return code means fail
402    if out["retcode"]:
403        ret = out["stderr"].strip()
404    else:
405        ret = out["stdout"].strip()
406
407    return ret
408
409
410def zero(protocol=None, service_address=None):
411    """
412
413    Zero the packet, byte and rate counters in a service or all services.
414
415    CLI Example:
416
417    .. code-block:: bash
418
419        salt '*' lvs.zero
420    """
421
422    if service_address:
423        cmd = "{} -Z {}".format(
424            __detect_os(),
425            _build_cmd(protocol=protocol, service_address=service_address),
426        )
427    else:
428        cmd = "{} -Z".format(__detect_os())
429    out = __salt__["cmd.run_all"](cmd, python_shell=False)
430
431    # A non-zero return code means fail
432    if out["retcode"]:
433        ret = out["stderr"].strip()
434    else:
435        ret = True
436    return ret
437
438
439def check_service(protocol=None, service_address=None, **kwargs):
440    """
441
442    Check the virtual service exists.
443
444    CLI Example:
445
446    .. code-block:: bash
447
448        salt '*' lvs.check_service tcp 1.1.1.1:80
449    """
450
451    cmd = "{}".format(
452        _build_cmd(protocol=protocol, service_address=service_address, **kwargs)
453    )
454    # Exact match
455    if not kwargs:
456        cmd += " "
457
458    all_rules = get_rules()
459    out = all_rules.find(cmd)
460
461    if out != -1:
462        ret = True
463    else:
464        ret = "Error: service not exists"
465    return ret
466
467
468def check_server(protocol=None, service_address=None, server_address=None, **kwargs):
469    """
470
471    Check the real server exists in the specified service.
472
473    CLI Example:
474
475    .. code-block:: bash
476
477         salt '*' lvs.check_server tcp 1.1.1.1:80 192.168.0.11:8080
478    """
479
480    cmd = "{}".format(
481        _build_cmd(
482            protocol=protocol,
483            service_address=service_address,
484            server_address=server_address,
485            **kwargs
486        )
487    )
488    # Exact match
489    if not kwargs:
490        cmd += " "
491
492    all_rules = get_rules()
493    out = all_rules.find(cmd)
494
495    if out != -1:
496        ret = True
497    else:
498        ret = "Error: server not exists"
499    return ret
500