Merge pull request #1178 from ogayot/FR-1953

Expose endpoint to get list of UA activable services

subiquity/client/controllers/ubuntu_advantage.py

@@ -15,22 +15,17 @@
 """ Module that defines the client-side controller class for Ubuntu Advantage.
 """
 
+import asyncio
 import logging
-import os
+from typing import Optional
 
 from subiquitycore.async_helpers import schedule_task
 
 from subiquity.client.controller import SubiquityTuiController
-from subiquity.common.ubuntu_advantage import (
-    InvalidUATokenError,
-    ExpiredUATokenError,
-    CheckSubscriptionError,
-    UAInterface,
-    UAInterfaceStrategy,
-    MockedUAInterfaceStrategy,
-    UAClientUAInterfaceStrategy,
-)
-from subiquity.common.types import UbuntuAdvantageInfo
+from subiquity.common.types import (
+    UbuntuAdvantageInfo,
+    UbuntuAdvantageCheckTokenStatus as TokenStatus,
+    )
 from subiquity.ui.views.ubuntu_advantage import UbuntuAdvantageView
 
 from subiquitycore.lsb_release import lsb_release
@@ -44,19 +39,10 @@ class UbuntuAdvantageController(SubiquityTuiController):
 
     endpoint_name = "ubuntu_advantage"
 
-    def __init__(self, app):
-        """ Initializer for client-side UA controller. """
-        strategy: UAInterfaceStrategy
-        if app.opts.dry_run:
-            strategy = MockedUAInterfaceStrategy(scale_factor=app.scale_factor)
-        else:
-            # Make sure we execute `$PYTHON "$SNAP/usr/bin/ubuntu-advantage"`.
-            executable = (
-                os.environ["PYTHON"],
-                os.path.join(os.environ["SNAP"], "usr/bin/ubuntu-advantage"),
-            )
-            strategy = UAClientUAInterfaceStrategy(executable=executable)
-        self.ua_interface = UAInterface(strategy)
+    def __init__(self, app) -> None:
+        """ Initializer for the client-side UA controller. """
+        self._check_task: Optional[asyncio.Future] = None
+
         super().__init__(app)
 
     async def make_ui(self) -> UbuntuAdvantageView:
@@ -77,21 +63,19 @@ def run_answers(self) -> None:
     def check_token(self, token: str):
         """ Asynchronously check the token passed as an argument. """
         async def inner() -> None:
-            try:
-                svcs = await \
-                        self.ua_interface.get_activable_services(token=token)
-            except InvalidUATokenError:
+            answer = await self.endpoint.check_token.GET(token)
+            if answer.status == TokenStatus.INVALID_TOKEN:
                 if isinstance(self.ui.body, UbuntuAdvantageView):
                     self.ui.body.show_invalid_token()
-            except ExpiredUATokenError:
+            elif answer.status == TokenStatus.EXPIRED_TOKEN:
                 if isinstance(self.ui.body, UbuntuAdvantageView):
                     self.ui.body.show_expired_token()
-            except CheckSubscriptionError:
+            elif answer.status == TokenStatus.UNKNOWN_ERROR:
                 if isinstance(self.ui.body, UbuntuAdvantageView):
                     self.ui.body.show_unknown_error()
             else:
                 if isinstance(self.ui.body, UbuntuAdvantageView):
-                    self.ui.body.show_activable_services(svcs)
+                    self.ui.body.show_activable_services(answer.services)
 
         self._check_task = schedule_task(inner())
 

subiquity/common/apidef.py

@@ -49,6 +49,7 @@
     StorageResponseV2,
     TimeZoneInfo,
     UbuntuAdvantageInfo,
+    UbuntuAdvantageCheckTokenAnswer,
     WLANSupportInstallState,
     ZdevInfo,
     WSLConfigurationBase,
@@ -321,6 +322,10 @@ def POST(data: Payload[UbuntuAdvantageInfo]) -> None: ...
         class skip:
             def POST() -> None: ...
 
+        class check_token:
+            def GET(token: Payload[str]) \
+                    -> UbuntuAdvantageCheckTokenAnswer: ...
+
 
 class LinkAction(enum.Enum):
     NEW = enum.auto()

subiquity/common/types.py

@@ -394,6 +394,26 @@ class UbuntuAdvantageInfo:
     token: str = attr.ib(repr=False)
 
 
+class UbuntuAdvantageCheckTokenStatus(enum.Enum):
+    VALID_TOKEN = enum.auto()
+    INVALID_TOKEN = enum.auto()
+    EXPIRED_TOKEN = enum.auto()
+    UNKNOWN_ERROR = enum.auto()
+
+
+@attr.s(auto_attribs=True)
+class UbuntuAdvantageService:
+    name: str
+    description: str
+
+
+@attr.s(auto_attribs=True)
+class UbuntuAdvantageCheckTokenAnswer:
+    status: UbuntuAdvantageCheckTokenStatus
+
+    services: Optional[List[UbuntuAdvantageService]]
+
+
 class ShutdownMode(enum.Enum):
     REBOOT = enum.auto()
     POWEROFF = enum.auto()

subiquity/server/controllers/ubuntu_advantage.py

@@ -15,9 +15,23 @@
 """ Module defining the server-side controller class for Ubuntu Advantage. """
 
 import logging
+import os
 
 from subiquity.common.apidef import API
-from subiquity.common.types import UbuntuAdvantageInfo
+from subiquity.common.types import (
+    UbuntuAdvantageInfo,
+    UbuntuAdvantageCheckTokenAnswer,
+    UbuntuAdvantageCheckTokenStatus,
+)
+from subiquity.server.ubuntu_advantage import (
+    InvalidUATokenError,
+    ExpiredUATokenError,
+    CheckSubscriptionError,
+    UAInterface,
+    UAInterfaceStrategy,
+    MockedUAInterfaceStrategy,
+    UAClientUAInterfaceStrategy,
+)
 from subiquity.server.controller import SubiquityController
 
 log = logging.getLogger("subiquity.server.controllers.ubuntu_advantage")
@@ -47,6 +61,21 @@ class UbuntuAdvantageController(SubiquityController):
         },
     }
 
+    def __init__(self, app) -> None:
+        """ Initializer for server-side UA controller. """
+        strategy: UAInterfaceStrategy
+        if app.opts.dry_run:
+            strategy = MockedUAInterfaceStrategy(scale_factor=app.scale_factor)
+        else:
+            # Make sure we execute `$PYTHON "$SNAP/usr/bin/ubuntu-advantage"`.
+            executable = (
+                os.environ["PYTHON"],
+                os.path.join(os.environ["SNAP"], "usr/bin/ubuntu-advantage"),
+            )
+            strategy = UAClientUAInterfaceStrategy(executable=executable)
+        self.ua_interface = UAInterface(strategy)
+        super().__init__(app)
+
     def load_autoinstall_data(self, data: dict) -> None:
         """ Load autoinstall data and update the model. """
         if data is None:
@@ -88,3 +117,25 @@ async def skip_POST(self) -> None:
         """ When running on a non-LTS release, we want to call this so we can
         skip the screen on the client side. """
         await self.configured()
+
+    async def check_token_GET(self, token: str) \
+            -> UbuntuAdvantageCheckTokenAnswer:
+        """ Handle a GET request asking whether the contract token is valid or
+        not. If it is valid, we provide the list of activable services
+        associated with the subscription.
+        """
+        services = None
+        try:
+            services = await \
+                    self.ua_interface.get_activable_services(token=token)
+        except InvalidUATokenError:
+            status = UbuntuAdvantageCheckTokenStatus.INVALID_TOKEN
+        except ExpiredUATokenError:
+            status = UbuntuAdvantageCheckTokenStatus.EXPIRED_TOKEN
+        except CheckSubscriptionError:
+            status = UbuntuAdvantageCheckTokenStatus.UNKNOWN_ERROR
+        else:
+            status = UbuntuAdvantageCheckTokenStatus.VALID_TOKEN
+
+        return UbuntuAdvantageCheckTokenAnswer(status=status,
+                                               services=services)

subiquity/common/tests/test_ubuntu_advantage.py → subiquity/server/tests/test_ubuntu_advantage.py

@@ -15,9 +15,10 @@
 
 from subprocess import CalledProcessError, CompletedProcess
 import unittest
-from unittest.mock import patch
+from unittest.mock import patch, AsyncMock
 
-from subiquity.common.ubuntu_advantage import (
+from subiquity.common.types import UbuntuAdvantageService
+from subiquity.server.ubuntu_advantage import (
     InvalidUATokenError,
     ExpiredUATokenError,
     CheckSubscriptionError,
@@ -56,7 +57,7 @@ def test_query_info_valid(self):
 
 
 class TestUAClientUAInterfaceStrategy(unittest.TestCase):
-    arun_command = "subiquity.common.ubuntu_advantage.utils.arun_command"
+    arun_command = "subiquity.server.ubuntu_advantage.utils.arun_command"
 
     def test_init(self):
         # Default initializer.
@@ -141,10 +142,65 @@ def test_mocked_get_activable_services(self):
             run_coro(interface.get_activable_services(token="xpiredToken"))
 
         # Other tokens are considered valid in dry-run mode.
+        run_coro(interface.get_activable_services(token="validToken"))
+
+    def test_get_activable_services(self):
+        # We use the standard strategy but don't actually run it
+        strategy = UAClientUAInterfaceStrategy()
+        interface = UAInterface(strategy)
+
+        subscription = {
+            "expires": "2035-12-31T00:00:00+00:00",
+            "services": [
+                {
+                    "name": "cis",
+                    "description": "Center for Internet Security Audit Tools",
+                    "entitled": "no",
+                    "auto_enabled": "no",
+                    "available": "yes"
+                },
+                {
+                    "name": "esm-apps",
+                    "description":
+                        "UA Apps: Extended Security Maintenance (ESM)",
+                    "entitled": "yes",
+                    "auto_enabled": "yes",
+                    "available": "no"
+                },
+                {
+                    "name": "esm-infra",
+                    "description":
+                        "UA Infra: Extended Security Maintenance (ESM)",
+                    "entitled": "yes",
+                    "auto_enabled": "yes",
+                    "available": "yes"
+                },
+                {
+                    "name": "fips",
+                    "description": "NIST-certified core packages",
+                    "entitled": "yes",
+                    "auto_enabled": "no",
+                    "available": "yes"
+                },
+            ]
+        }
+        interface.get_subscription = AsyncMock(return_value=subscription)
         services = run_coro(
-                interface.get_activable_services(token="validToken"))
-        for service in services:
-            self.assertIn("name", service)
-            self.assertIn("description", service)
-            self.assertEqual(service["available"], "yes")
-            self.assertEqual(service["entitled"], "yes")
+                interface.get_activable_services(token="XXX"))
+
+        self.assertIn(UbuntuAdvantageService(
+            name="esm-infra",
+            description="UA Infra: Extended Security Maintenance (ESM)",
+        ), services)
+        self.assertIn(UbuntuAdvantageService(
+            name="fips",
+            description="NIST-certified core packages",
+        ), services)
+        self.assertNotIn(UbuntuAdvantageService(
+            name="esm-apps",
+            description="UA Apps: Extended Security Maintenance (ESM)",
+        ), services)
+        self.assertNotIn(UbuntuAdvantageService(
+            name="cis",
+            description="Center for Internet Security Audit Tools",
+        ), services)

subiquity/common/ubuntu_advantage.py → subiquity/server/ubuntu_advantage.py

@@ -23,10 +23,11 @@
 from typing import List, Sequence, Union
 import asyncio
 
+from subiquity.common.types import UbuntuAdvantageService
 from subiquitycore import utils
 
 
-log = logging.getLogger("subiquitycore.common.ubuntu_advantage")
+log = logging.getLogger("subiquity.server.ubuntu_advantage")
 
 
 class InvalidUATokenError(Exception):
@@ -147,7 +148,8 @@ async def get_subscription(self, token: str) -> dict:
         """ Return a dictionary containing the subscription information. """
         return await self.strategy.query_info(token)
 
-    async def get_activable_services(self, token: str) -> List[dict]:
+    async def get_activable_services(self, token: str) \
+            -> List[UbuntuAdvantageService]:
         """ Return a list of activable services (i.e. services that are
         entitled to the subscription and available on the current hardware).
         """
@@ -165,4 +167,16 @@ def is_activable_service(service: dict) -> bool:
             return service["available"] == "yes" \
                and service["entitled"] == "yes"
 
-        return [svc for svc in info["services"] if is_activable_service(svc)]
+        def service_from_dict(service: dict) -> UbuntuAdvantageService:
+            return UbuntuAdvantageService(
+               name=service["name"],
+               description=service["description"],
+            )
+
+        activable_services: List[UbuntuAdvantageService] = []
+
+        for service in info["services"]:
+            if not is_activable_service(service):
+                continue
+            activable_services.append(service_from_dict(service))
+        return activable_services

subiquity/ui/views/ubuntu_advantage.py

@@ -25,6 +25,7 @@
     Widget,
     )
 
+from subiquity.common.types import UbuntuAdvantageService as UAService
 from subiquitycore.view import BaseView
 from subiquitycore.ui.buttons import (
     back_btn,
@@ -187,7 +188,7 @@ def show_unknown_error(self) -> None:
         self.remove_overlay()
         self.show_stretchy_overlay(ContinueAnywayWidget(self))
 
-    def show_activable_services(self, services: List[dict]) -> None:
+    def show_activable_services(self, services: List[UAService]) -> None:
         """ Display an overlay with the list of services that can be enabled
         via Ubuntu Advantage subscription. After the user confirms, we will
         quit the current view and move on. """
@@ -198,7 +199,7 @@ def show_activable_services(self, services: List[dict]) -> None:
 class ShowServicesWidget(Stretchy):
     """ Widget to show the activable services for UA subscription. """
     def __init__(self, parent: UbuntuAdvantageView,
-                 services: List[dict]) -> None:
+                 services: List[UAService]) -> None:
         """ Initializes the widget by including the list of services as a
         bullet-point list. """
         self.parent = parent
@@ -212,7 +213,7 @@ def __init__(self, parent: UbuntuAdvantageView,
         widgets: List[Widget] = [
             Text(header),
             Text(""),
-            Pile([Text(f"* {svc['description']}") for svc in services]),
+            Pile([Text(f"* {svc.description}") for svc in services]),
             Text(""),
             Text("Once the installation has finished, you can enable these "
                  "services using the `ua` command-line tool."),