1# Copyright 2020 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5import json
6import logging
7import os
8import shutil
9
10from autotest_lib.client.common_lib import utils
11
12_DEFAULT_RUN = utils.run
13
14# Directory for preloaded DLCs that may be packaged with the OS. Preloaded
15# DLCs can be installed without needing to go through update_engine.
16_PRELOAD_DIR = '/mnt/stateful_partition/var_overlay/cache/dlc-images/'
17
18class DLCUtil(object):
19    """
20    Wrapper around dlcservice_util for tests that require DLC management.
21
22    For dlc_util to work seamlessly in both client and server tests, we need
23    those tests to define the _run() function:
24    server side: host.run
25    client side: utils.run
26
27    """
28    _DLCSERVICE_UTIL_CMD = "dlcservice_util"
29    _SAMPLE_DLC_ID = "sample-dlc"
30
31    def __init__(self, run_func=_DEFAULT_RUN):
32        """
33        Initialize this class with _run() function.
34
35        @param run_func: The function to use to run commands on the client.
36                         Defaults for use by client tests, but can be
37                         overwritten to run remotely from a server test.
38
39        """
40        self._run = run_func
41
42
43    def list(self):
44        """
45        List DLCs that are installed on the DUT and additional information
46        about them such as name, version, root_mount, and size. The output is
47        a dictionary containing the result of dlcservice_util --list, which
48        looks like:
49        {
50           "sample-dlc": [ {
51              "fs-type": "squashfs",
52              "id": "sample-dlc",
53              "image_type": "dlc",
54              "manifest": "/opt/google/dlc/sample-dlc/package/imageloader.json",
55              "name": "Sample DLC",
56              "package": "package",
57              "preallocated_size": "4194304",
58              "root_mount": "/run/imageloader/sample-dlc/package",
59              "size": "53248",
60              "version": "1.0.0-r10"
61           } ]
62        }
63
64        @return Dictionary containing information about the installed DLCs,
65                whose keys are the DLC IDs.
66
67        """
68        status = self._run([self._DLCSERVICE_UTIL_CMD, '--list'])
69        logging.info(status)
70
71        return json.loads(status.stdout)
72
73
74    def install(self, dlc_id, omaha_url, timeout=900):
75        """
76        Install a DLC on the stateful partition.
77
78        @param dlc_id: The id of the DLC to install.
79        @param omaha_url: The Omaha URL to send the install request to.
80
81        """
82        cmd = [self._DLCSERVICE_UTIL_CMD, '--install', '--id=%s' % dlc_id,
83               '--omaha_url=%s' % omaha_url]
84        self._run(cmd, timeout=timeout)
85
86
87    def uninstall(self, dlc_id, ignore_status=False):
88        """
89        Uninstall a DLC. The DLC will remain on the DUT in an un-mounted state
90        and can be reinstalled without requiring an install request. To
91        completely remove a DLC, purge it instead.
92
93        @param dlc_id: The id of the DLC to uninstall.
94        @param ignore_status: Whether or not to ignore the return status when
95                              running the uninstall command.
96
97        """
98        cmd = [self._DLCSERVICE_UTIL_CMD, '--uninstall', '--id=%s' % dlc_id]
99        self._run(cmd, ignore_status=ignore_status)
100
101
102    def purge(self, dlc_id, ignore_status=False):
103        """
104        Purge a DLC. This will completely remove the DLC from the device.
105
106        @param dlc_id: The id of the DLC to purge.
107        @param ignore_status: Whether or not to ignore the return status when
108                              running the purge command.
109
110        """
111        cmd = [self._DLCSERVICE_UTIL_CMD, '--purge', '--id=%s' % dlc_id]
112        self._run(cmd, ignore_status=ignore_status)
113
114
115    def remove_preloaded(self, dlc_id):
116        """
117        Remove a DLC from the preload directory. DLCs in this directory can be
118        preloaded, meaning they can be installed without needing to download
119        and install them using update_engine. It is hard to differentiate
120        preloading a DLC and installing a DLC after a successful AU. After
121        AU, updated DLCs should be installed but not yet mounted. In both
122        cases, calling dlcservice_util --install should result in the DLC
123        being installed without going through Omaha/Nebraska. Clearing the
124        preload directory for a DLC allows us to verify it was updated,
125        by ruling out the possibility that it was preloaded instead.
126
127        @param dlc_id: Wipe preload directory for the DLC with this id.
128
129        """
130        preload_dir = os.path.join(_PRELOAD_DIR, dlc_id)
131        if os.path.exists(preload_dir) and os.path.isdir(preload_dir):
132            shutil.rmtree(preload_dir)
133
134
135    def is_installed(self, dlc_id):
136        """
137        Check if a DLC is installed.
138
139        @param dlc_id: The id of the DLC to check.
140
141        @return True if the DLC is installed, False if it's not.
142
143        """
144        return dlc_id in self.list().keys()
145