Merge pull request #16 from GoogleCloudPlatform/develop

0.3-beta

.bumpversion.cfg

@@ -0,0 +1,17 @@
+[bumpversion]
+current_version = 0.3-beta
+parse = (?P<major>\d+)\.(?P<minor>\d+)(-(?P<release>.*))?
+message = Bump version: {current_version} -> {new_version}
+serialize = 
+	{major}.{minor}-{release}
+	{major}.{minor}
+
+[bumpversion:part:release]
+optional_value = release
+values = 
+	beta
+	release
+
+[bumpversion:file:gce_rescue/config.py]
+search = VERSION = '{current_version}'
+replace = VERSION = '{new_version}'

.gitignore

@@ -22,4 +22,4 @@ poetry.lock
 *.DS_Store
 .vscode/
 .idea
-debian.log
+*.log

Makefile

@@ -7,6 +7,9 @@ test:
 requirements: requirements.txt
 	python3 -m pip install -r requirements.txt
 
+bumpversion:
+	pipenv run bumpversion --commit minor
+
 build: setup.py
 	python3 ./setup.py bdist_wheel sdist
 

Pipfile

@@ -8,6 +8,7 @@ absl-py = ">=1.2.0"
 google-api-python-client = "*"
 google-auth = "*"
 pylint = "*"
+bumpversion = "*"
 
 [dev-packages]
 bump2version = "*"

gce_rescue/bin/rescue.py

@@ -21,7 +21,7 @@
 
 from gce_rescue.config import process_args, set_configs
 from gce_rescue import messages
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 from gce_rescue.tasks.actions import call_tasks
 from gce_rescue.utils import read_input, set_logging
 
@@ -75,4 +75,4 @@ def main():
 
 
 if __name__ == '__main__':
-  main()
+  main()

gce_rescue/config.py

@@ -19,7 +19,10 @@
 
 dirname = os.path.dirname(__file__)
 
+VERSION = '0.3-beta'
+
 config = {
+  'version': VERSION,
   'debug': False,
   'startup-script-file': os.path.join(dirname, 'startup-script.txt'),
   'source_guests': {
@@ -41,8 +44,8 @@ def get_config(key):
 
 def process_args():
   """ Print usage options. """
-  parser = argparse.ArgumentParser(description='GCE Rescue v0.0.2-1 - Set/Reset\
-    GCE instances to boot in rescue mode.')
+  parser = argparse.ArgumentParser(description=f'GCE Rescue v{VERSION} - \
+    Set/Reset GCE instances to boot in rescue mode.')
   parser.add_argument('-p', '--project',
                       help='The project-id that has the instance.')
   parser.add_argument('-z', '--zone', help='Zone where the instance \

gce_rescue/rescue.py → gce_rescue/gce.py

@@ -14,27 +14,87 @@
 
 """ Initilization Instance() with VM information. """
 
+from googleapiclient.discovery import Resource
+
 from dataclasses import dataclass, field
 from typing import Dict, List, Union
-
+from time import time
 from gce_rescue.tasks.backup import backup_metadata_items
 from gce_rescue.tasks.disks import list_disk
 from gce_rescue.tasks.pre_validations import Validations
-from gce_rescue.utils import (
-  validate_instance_mode,
-  guess_guest,
-  get_instance_info
-)
-import googleapiclient.discovery
+from gce_rescue.config import get_config
+
+
+def get_instance_info(
+  compute: Resource,
+  name: str,
+  project_data: Dict[str, str]
+) -> Dict:
+  """Set Dictionary with complete data from instances().get() from the instance.
+  https://cloud.google.com/compute/docs/reference/rest/v1/instances/get
+  Attributes:
+    compute: obj, API Object
+    instance: str, Instace name
+    project_data: dict, Dictionary containing project and zone keys to be
+      unpacked when calling the API.
+  """
+  return compute.instances().get(
+      **project_data,
+      instance = name).execute()
+
+
+def guess_guest(data: Dict) -> str:
+  """Determined which Guest OS Family is being used and select a
+  different OS for recovery disk.
+     Default: projects/debian-cloud/global/images/family/debian-11"""
+
+  guests = get_config('source_guests')
+  for disk in data['disks']:
+    if disk['boot']:
+      if 'architecture' in disk:
+        arch = disk['architecture'].lower()
+      else:
+        arch = 'x86_64'
+      guest_default = guests[arch][0]
+      guest_name = guest_default.split('/')[-1]
+      for lic in disk['licenses']:
+        if guest_name in lic:
+          guest_default = guests[arch][1]
+  return guest_default
+
+
+def validate_instance_mode(data: Dict) -> Dict:
+  """Validate if the instance is already configured as rescue mode."""
+
+  result = {
+      'rescue-mode': False,
+      'ts': generate_ts()
+  }
+  if 'metadata' in data and  'items' in data['metadata']:
+    metadata = data['metadata']
+    for item in metadata['items']:
+      if item['key'] == 'rescue-mode':
+        result = {
+          'rescue-mode': True,
+          'ts': item['value']
+        }
+
+  return result
+
+def generate_ts() -> int:
+  """Get the current timestamp to be used as unique ID
+  during this execution."""
+  return int(time())
+
 
 @dataclass
-class Instance:
+class Instance(Resource):
   """Initialize instance."""
   zone: str
   name: str
   project: str = None
   test_mode: bool = field(default_factory=False)
-  compute: googleapiclient.discovery.Resource = field(init=False)
+  compute: Resource = field(init=False)
   data: Dict[str, Union[str, int]] = field(init=False)
   ts: int = field(init=False)
   _status: str = ''

gce_rescue/rescue_test.py → gce_rescue/gce_test.py

@@ -15,7 +15,7 @@
 """Test code for rescue.py."""
 
 from absl.testing import absltest
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 from gce_rescue.test.mocks import (
   mock_api_object,
   MOCK_TEST_VM,

gce_rescue/messages.py

@@ -14,7 +14,7 @@
 
 """ List of messages to inform and educate the user. """
 
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 
 def tip_connect_ssh(vm: Instance) -> str:
   return (f'└── Your instance is READY! You can now connect your instance '

gce_rescue/messages_test.py

@@ -16,7 +16,7 @@
 
 from absl.testing import absltest
 from gce_rescue import messages
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 from gce_rescue.test.mocks import mock_api_object, MOCK_TEST_VM
 
 

gce_rescue/tasks/actions.py

@@ -14,10 +14,10 @@
 
 """ List of ordered tasks to be executed when set/reset VM rescue mode. """
 
-from typing import Dict
+from typing import List
 import logging
 
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 from gce_rescue.tasks.disks import (
   config_rescue_disks,
   restore_original_disk,
@@ -35,7 +35,7 @@
 
 _logger = logging.getLogger(__name__)
 
-def _list_tasks(vm: Instance, action: str) -> Dict:
+def _list_tasks(vm: Instance, action: str) -> List:
   """ List tasks, by order, per operation
     operations (str):
       1. set_rescue_mode
@@ -113,7 +113,7 @@ def _list_tasks(vm: Instance, action: str) -> Dict:
 
   if action not in all_tasks:
     _logger.info(f'Unable to find "{action}".')
-    raise Exception(ValueError)
+    raise ValueError()
   return all_tasks[action]
 
 

gce_rescue/tasks/backup.py

@@ -15,7 +15,7 @@
 """ Different operations to guarantee VM disks backup, before performing
     any modifications."""
 
-from gce_rescue.utils import wait_for_operation
+from gce_rescue.tasks.keeper import wait_for_operation
 from typing import Dict, List
 import logging
 

gce_rescue/tasks/backup_test.py

@@ -17,7 +17,7 @@
 from absl.testing import absltest
 from gce_rescue.tasks import backup
 from gce_rescue.test.mocks import mock_api_object, MOCK_TEST_VM
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 
 
 class BackupTest(absltest.TestCase):

gce_rescue/tasks/disks.py

@@ -19,9 +19,9 @@
 
 import googleapiclient.errors
 
-from gce_rescue.utils import wait_for_operation
+from gce_rescue.tasks.keeper import wait_for_operation
 from gce_rescue.tasks.backup import  backup
-from gce_rescue.multitasks import Handler
+from gce_rescue.utils import ThreadHandler as Handler
 
 _logger = logging.getLogger(__name__)
 

gce_rescue/tasks/disks_test.py

@@ -17,7 +17,7 @@
 from absl.testing import absltest
 from absl import logging
 
-from gce_rescue.rescue import Instance
+from gce_rescue.gce import Instance
 from gce_rescue.tasks import disks
 from gce_rescue.test.mocks import (
   mock_api_object,

gce_rescue/tasks/keeper.py

@@ -0,0 +1,66 @@
+# Copyright 2021 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=broad-exception-raised
+
+"""keeper that the progress of the tasks. """
+
+import googleapiclient.discovery
+from time import sleep
+from typing import Dict
+import logging
+import json
+
+_logger = logging.getLogger(__name__)
+
+
+def wait_for_operation(
+  instance_obj: googleapiclient.discovery.Resource,
+  oper: Dict
+) -> Dict:
+  """ Creating poll to wait the operation to finish. """
+
+  while True:
+    if oper['status'] == 'DONE':
+      _logger.info('done.')
+      if 'error' in oper:
+        raise Exception(oper['error'])
+      return oper
+
+    oper = instance_obj.compute.zoneOperations().get(
+      **instance_obj.project_data,
+      operation = oper['name']).execute()
+    sleep(1)
+
+def wait_for_os_boot(vm: googleapiclient.discovery.Resource) -> bool:
+  """Wait guest OS to complete the boot proccess."""
+
+  timeout = 60
+  wait_time = 2
+  end_string = f'END:{vm.ts}'
+  _logger.info('Waiting startup-script to complete.')
+  while True:
+    result = vm.compute.instances().getSerialPortOutput(
+      **vm.project_data,
+      instance = vm.name
+    ).execute()
+
+    if end_string in json.dumps(result):
+      _logger.info('startup-script has ended.')
+      return True
+
+    sleep(wait_time)
+    timeout -= wait_time
+    if not timeout:
+      return False