1#!/usr/bin/env python
2#
3# Copyright 2016 - The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Config manager.
18
19Three protobuf messages are defined in
20   driver/internal/config/proto/internal_config.proto
21   driver/internal/config/proto/user_config.proto
22
23Internal config file     User config file
24      |                         |
25      v                         v
26  InternalConfig           UserConfig
27  (proto message)        (proto message)
28        |                     |
29        |                     |
30        |->   AcloudConfig  <-|
31
32At runtime, AcloudConfigManager performs the following steps.
33- Load driver config file into a InternalConfig message instance.
34- Load user config file into a UserConfig message instance.
35- Create AcloudConfig using InternalConfig and UserConfig.
36
37TODO(fdeng):
38  1. Add support for override configs with command line args.
39  2. Scan all configs to find the right config for given branch and build_id.
40     Raise an error if the given build_id is smaller than min_build_id
41     only applies to release build id.
42     Raise an error if the branch is not supported.
43
44"""
45
46import logging
47import os
48
49from acloud.internal.proto import internal_config_pb2
50from acloud.internal.proto import user_config_pb2
51from acloud.public import errors
52from google.protobuf import text_format
53_CONFIG_DATA_PATH = os.path.join(
54    os.path.dirname(os.path.abspath(__file__)), "data")
55
56logger = logging.getLogger(__name__)
57
58
59class AcloudConfig(object):
60    """A class that holds all configurations for acloud."""
61
62    REQUIRED_FIELD = [
63        "project", "zone", "machine_type", "network", "storage_bucket_name",
64        "min_machine_size", "disk_image_name", "disk_image_mime_type"
65    ]
66
67    def __init__(self, usr_cfg, internal_cfg):
68        """Initialize.
69
70        Args:
71            usr_cfg: A protobuf object that holds the user configurations.
72            internal_cfg: A protobuf object that holds internal configurations.
73        """
74        self.service_account_name = usr_cfg.service_account_name
75        self.service_account_private_key_path = (
76            usr_cfg.service_account_private_key_path)
77        self.creds_cache_file = internal_cfg.creds_cache_file
78        self.user_agent = internal_cfg.user_agent
79        self.client_id = usr_cfg.client_id
80        self.client_secret = usr_cfg.client_secret
81
82        self.project = usr_cfg.project
83        self.zone = usr_cfg.zone
84        self.machine_type = (usr_cfg.machine_type or
85                             internal_cfg.default_usr_cfg.machine_type)
86        self.network = (usr_cfg.network or
87                        internal_cfg.default_usr_cfg.network)
88        self.ssh_private_key_path = usr_cfg.ssh_private_key_path
89        self.ssh_public_key_path = usr_cfg.ssh_public_key_path
90        self.storage_bucket_name = usr_cfg.storage_bucket_name
91        self.metadata_variable = {
92            key: val
93            for key, val in
94            internal_cfg.default_usr_cfg.metadata_variable.iteritems()
95        }
96        self.metadata_variable.update(usr_cfg.metadata_variable)
97
98        self.device_resolution_map = {
99            device: resolution
100            for device, resolution in
101            internal_cfg.device_resolution_map.iteritems()
102        }
103        self.device_default_orientation_map = {
104            device: orientation
105            for device, orientation in
106            internal_cfg.device_default_orientation_map.iteritems()
107        }
108        self.no_project_access_msg_map = {
109            project: msg for project, msg
110            in internal_cfg.no_project_access_msg_map.iteritems()}
111        self.min_machine_size = internal_cfg.min_machine_size
112        self.disk_image_name = internal_cfg.disk_image_name
113        self.disk_image_mime_type = internal_cfg.disk_image_mime_type
114        self.disk_image_extension = internal_cfg.disk_image_extension
115        self.disk_raw_image_name = internal_cfg.disk_raw_image_name
116        self.disk_raw_image_extension = internal_cfg.disk_raw_image_extension
117        self.valid_branch_and_min_build_id = {
118            branch: min_build_id
119            for branch, min_build_id in
120            internal_cfg.valid_branch_and_min_build_id.iteritems()
121        }
122        self.precreated_data_image_map = {
123            size_gb: image_name
124            for size_gb, image_name in
125            internal_cfg.precreated_data_image.iteritems()
126        }
127        self.extra_data_disk_size_gb = (
128            usr_cfg.extra_data_disk_size_gb or
129            internal_cfg.default_usr_cfg.extra_data_disk_size_gb)
130        if self.extra_data_disk_size_gb > 0:
131            if "cfg_sta_persistent_data_device" not in usr_cfg.metadata_variable:
132                # If user did not set it explicity, use default.
133                self.metadata_variable["cfg_sta_persistent_data_device"] = (
134                    internal_cfg.default_extra_data_disk_device)
135            if "cfg_sta_ephemeral_data_size_mb" in usr_cfg.metadata_variable:
136                raise errors.ConfigError(
137                    "The following settings can't be set at the same time: "
138                    "extra_data_disk_size_gb and"
139                    "metadata variable cfg_sta_ephemeral_data_size_mb.")
140            if "cfg_sta_ephemeral_data_size_mb" in self.metadata_variable:
141                del self.metadata_variable["cfg_sta_ephemeral_data_size_mb"]
142
143        # Fields that can be overriden by args
144        self.orientation = usr_cfg.orientation
145        self.resolution = usr_cfg.resolution
146
147        # Verify validity of configurations.
148        self.Verify()
149
150    def OverrideWithArgs(self, parsed_args):
151        """Override configuration values with args passed in from cmd line.
152
153        Args:
154            parsed_args: Args parsed from command line.
155        """
156        if parsed_args.which == "create" and parsed_args.spec:
157            if not self.resolution:
158                self.resolution = self.device_resolution_map.get(
159                    parsed_args.spec, "")
160            if not self.orientation:
161                self.orientation = self.device_default_orientation_map.get(
162                    parsed_args.spec, "")
163        if parsed_args.email:
164            self.service_account_name = parsed_args.email
165
166    def Verify(self):
167        """Verify configuration fields."""
168        missing = [f for f in self.REQUIRED_FIELD if not getattr(self, f)]
169        if missing:
170            raise errors.ConfigError(
171                "Missing required configuration fields: %s" % missing)
172        if (self.extra_data_disk_size_gb and self.extra_data_disk_size_gb
173                not in self.precreated_data_image_map):
174            raise errors.ConfigError(
175                "Supported extra_data_disk_size_gb options(gb): %s, "
176                "invalid value: %d" % (self.precreated_data_image_map.keys(),
177                                       self.extra_data_disk_size_gb))
178
179
180class AcloudConfigManager(object):
181    """A class that loads configurations."""
182
183    _DEFAULT_INTERNAL_CONFIG_PATH = os.path.join(_CONFIG_DATA_PATH,
184                                                 "default.config")
185
186    def __init__(self,
187                 user_config_path,
188                 internal_config_path=_DEFAULT_INTERNAL_CONFIG_PATH):
189        """Initialize.
190
191        Args:
192            user_config_path: path to the user config.
193            internal_config_path: path to the internal conifg.
194        """
195        self._user_config_path = user_config_path
196        self._internal_config_path = internal_config_path
197
198    def Load(self):
199        """Load the configurations."""
200        internal_cfg = None
201        usr_cfg = None
202        try:
203            with open(self._internal_config_path) as config_file:
204                internal_cfg = self.LoadConfigFromProtocolBuffer(
205                    config_file, internal_config_pb2.InternalConfig)
206
207            with open(self._user_config_path, "r") as config_file:
208                usr_cfg = self.LoadConfigFromProtocolBuffer(
209                    config_file, user_config_pb2.UserConfig)
210        except OSError as e:
211            raise errors.ConfigError("Could not load config files: %s" %
212                                     str(e))
213        return AcloudConfig(usr_cfg, internal_cfg)
214
215    @staticmethod
216    def LoadConfigFromProtocolBuffer(config_file, message_type):
217        """Load config from a text-based protocol buffer file.
218
219        Args:
220            config_file: A python File object.
221            message_type: A proto message class.
222
223        Returns:
224            An instance of type "message_type" populated with data
225            from the file.
226        """
227        try:
228            config = message_type()
229            text_format.Merge(config_file.read(), config)
230            return config
231        except text_format.ParseError as e:
232            raise errors.ConfigError("Could not parse config: %s" % str(e))
233