xref: /aosp_15_r20/external/autotest/server/cros/network/ip_config_context_manager.py (revision 9c5db1993ded3edbeafc8092d69fe5de2ee02df7) 
1# Lint as: python2, python3
2# Copyright (c) 2021 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6import logging
7
8
9class IpConfigContextManager(object):
10    """Allows changes to IP configs on multiple host test devices which are
11    guaranteed to be reverted when the context is exited.
12    """
13
14    def bring_interface_up(self, host, dev_name):
15        """Bring a device interface up on the host. This interface will
16        automatically be brought down when the context is exited.
17
18        @param host Host Device to bring the interface up on.
19        @param dev_name String name of the device to bring up.
20
21        """
22        clear_command = 'sudo ip link set %s down' % dev_name
23        if host in self._iface_cleanup_dict:
24            self._iface_cleanup_dict[host].append(clear_command)
25        else:
26            self._iface_cleanup_dict[host] = [clear_command]
27        host.run('sudo ip link set %s up' % dev_name)
28
29    def add_ip_route(self, host, dest_ip, iface_name, via_ip=None):
30        """Add an ip route to the device. This route will be deleted when the
31        context is exited.
32
33        @param host Host Device to assign the ip route on.
34        @param dest_ip String destination ip address of the ip route.
35        @param iface_name String The local iface to route the traffic from.
36        @param via_ip String an optional ip address to route the traffic through.
37
38        """
39        via = "via %s " % via_ip if via_ip else ""
40        clear_command = 'sudo ip route del table 255 %s %sdev %s' % (
41                dest_ip, via, iface_name)
42        if host in self._ip_route_cleanup_dict:
43            self._ip_route_cleanup_dict[host].append(clear_command)
44        else:
45            self._ip_route_cleanup_dict[host] = [clear_command]
46        host.run('sudo ip route replace table 255 %s %sdev %s' %
47                 (dest_ip, via, iface_name))
48
49    def assign_ip_addr_to_iface(self, host, ip_addr, iface_name):
50        """Assign an ip address to an interface on the host. This address will be
51        deleted when the context is exited.
52
53        @param host Host Device to assign the ip address on.
54        @param ip_addr String ip address to assign.
55        @param iface_name String The interface to assign the ip address to.
56
57        """
58        clear_command = 'sudo ip addr del %s/24 dev %s' % (ip_addr, iface_name)
59        if host in self._ip_addr_cleanup_dict:
60            self._ip_addr_cleanup_dict[host].append(clear_command)
61        else:
62            self._ip_addr_cleanup_dict[host] = [clear_command]
63        host.run('sudo ip addr replace %s/24 dev %s' % (ip_addr, iface_name))
64
65    def __init__(self):
66        """Construct an IpConfigContextManager. This class uses dictionaries to
67        store the cleanup commands that must be run on various hosts when the
68        context is exited.
69        """
70        self._iface_cleanup_dict = dict()
71        self._ip_route_cleanup_dict = dict()
72        self._ip_addr_cleanup_dict = dict()
73
74    def __enter__(self):
75        return self
76
77    def __exit__(self, exc_type, exc_value, traceback):
78        logging.info('Cleaning up ip configs from test devices.')
79        for host in self._ip_route_cleanup_dict:
80            for command in self._ip_route_cleanup_dict[host]:
81                host.run(command)
82
83        for host in self._ip_addr_cleanup_dict:
84            for command in self._ip_addr_cleanup_dict[host]:
85                host.run(command)
86
87        for host in self._iface_cleanup_dict:
88            for command in self._iface_cleanup_dict[host]:
89                host.run(command)
90