1#!/usr/bin/env python3
2#
3#   Copyright 2020 - 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
17import logging
18
19from blueberry.tests.gd.cert.captures import SecurityCaptures
20from blueberry.tests.gd.cert.closable import Closable
21from blueberry.tests.gd.cert.closable import safeClose
22from blueberry.tests.gd.cert.event_stream import EventStream
23from blueberry.tests.gd.cert.matchers import SecurityMatchers
24from blueberry.tests.gd.cert.truth import assertThat
25from datetime import timedelta
26from google.protobuf import empty_pb2 as empty_proto
27from blueberry.facade.security.facade_pb2 import LeAuthRequirementsMessage
28from blueberry.facade.security.facade_pb2 import HelperMsgType
29
30
31class PyLeSecurity(Closable):
32    """
33        Abstraction for security tasks and GRPC calls
34    """
35
36    _ui_event_stream = None
37    _bond_event_stream = None
38    _helper_event_stream = None
39
40    def __init__(self, device):
41        logging.info("DUT: Init")
42        self._device = device
43        self._device.wait_channel_ready()
44        self._ui_event_stream = EventStream(self._device.security.FetchUiEvents(empty_proto.Empty()))
45        self._bond_event_stream = EventStream(self._device.security.FetchBondEvents(empty_proto.Empty()))
46        self._helper_event_stream = EventStream(self._device.security.FetchHelperEvents(empty_proto.Empty()))
47        self._advertising_callback_event_stream = EventStream(
48            self._device.security.FetchAdvertisingCallbackEvents(empty_proto.Empty()))
49
50    def get_ui_stream(self):
51        return self._ui_event_stream
52
53    def get_bond_stream(self):
54        return self._bond_event_stream
55
56    def get_advertising_callback_event_stream(self):
57        return self._advertising_callback_event_stream
58
59    def wait_for_ui_event_passkey(self, timeout=timedelta(seconds=3)):
60        display_passkey_capture = SecurityCaptures.DisplayPasskey()
61        assertThat(self._ui_event_stream).emits(display_passkey_capture, timeout=timeout)
62        return display_passkey_capture.get()
63
64    def wait_device_disconnect(self, address):
65        assertThat(self._helper_event_stream).emits(
66            SecurityMatchers.HelperMsg(HelperMsgType.DEVICE_DISCONNECTED, address))
67
68    def SetLeAuthRequirements(self, *args, **kwargs):
69        return self._device.security.SetLeAuthRequirements(LeAuthRequirementsMessage(*args, **kwargs))
70
71    def close(self):
72        if self._ui_event_stream is not None:
73            safeClose(self._ui_event_stream)
74        else:
75            logging.info("DUT: UI Event Stream is None!")
76
77        if self._bond_event_stream is not None:
78            safeClose(self._bond_event_stream)
79        else:
80            logging.info("DUT: Bond Event Stream is None!")
81
82        if self._helper_event_stream is not None:
83            safeClose(self._helper_event_stream)
84        else:
85            logging.info("DUT: Helper Event Stream is None!")
86
87        if self._advertising_callback_event_stream is not None:
88            safeClose(self._advertising_callback_event_stream)
89        else:
90            logging.info("DUT: Advertising Callback Event Stream is None!")
91