charm.py

#!/usr/bin/env python3
# Copyright 2020 Canonical Ltd.
#
# 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.

import sys

sys.path.append("lib")

from ops.charm import CharmBase, CharmEvents
from ops.framework import StoredState, EventBase, EventSource
from ops.main import main
from ops.model import (
    ActiveStatus,
    BlockedStatus,
    MaintenanceStatus,
    WaitingStatus,
    ModelError,
)

from charms.osm.sshproxy import SSHProxyCharm
import subprocess


class MagmaAGWProxyCharm(SSHProxyCharm):

    state = StoredState()

    def __init__(self, framework, key):
        super().__init__(framework, key)

        self.state.set_default(registered=False, ready=False)

        # Listen to charm events
        self.framework.observe(self.on.config_changed, self.on_config_changed)
        self.framework.observe(self.on.install, self.on_install)
        self.framework.observe(self.on.start, self.on_start)

        # Charm actions (primitives)
        self.framework.observe(self.on.add_net_action, self.on_add_net_action)
        self.framework.observe(self.on.add_gw_action, self.on_add_gw_action)
        self.framework.observe(self.on.reset_id_action, self.on_reset_id_action)
        self.framework.observe(self.on.add_hosts_action, self.on_add_hosts_action)
        self.framework.observe(
            self.on.restart_magma_action, self.on_restart_magma_action
        )
        self.framework.observe(self.on.del_gw_action, self.on_del_gw_action)
        self.framework.observe(self.on.reset_id_action, self.on_reset_id_action)
        self.framework.observe(
            self.on.add_test_subscriber_action, self.on_add_test_subscriber_action
        )

        # # Relation
        self.framework.observe(self.on.agw_relation_joined, self.agw_relation_joined)

    def agw_relation_joined(self, event):
        if self.unit.is_leader():
            if not self.state.ready:
                event.defer()
                return
            self.send_relation_data()

    def send_relation_data(self):
        relation = self.model.get_relation("agw")
        if relation is not None and self.state.mme_addr and self.state.magmagw:
            relation.data[self.unit]["mme-addr"] = self.state.mme_addr
            relation.data[self.unit]["magmagw"] = self.state.magmagw

    def on_config_changed(self, event):
        """Handle changes in configuration"""
        super().on_config_changed(event)

    def on_install(self, event):
        """Called when the charm is being installed"""
        super().on_install(event)

    def on_start(self, event):
        """Called when the charm is being started"""
        super().on_start(event)
        if not self.verify_credentials():
            event.defer()
            return
        proxy = self.get_ssh_proxy()
        ips, _ = proxy.run("hostname -I")
        ip_list = ips.split(" ")
        self.state.mme_addr = ip_list[0]
        self.state.magmagw = ip_list[1]

    # Magma AGW Action implementation
    def on_add_net_action(self, event):
        """Add AGW Network if needed"""
        if self.unit.is_leader():
            orch_ip = event.params["orch_ip"]
            orch_net = event.params["orch_net"]
            proxy = self.get_ssh_proxy()

            attempt = 0
            while attempt < 50:
                try:
                    stdout, stderr = proxy.run(
                        "/home/magma/addnet.py --orch_ip {} --orch_net {}".format(
                            orch_ip, orch_net
                        )
                    )
                    break
                except subprocess.CalledProcessError:
                    attempt += 1
                    import time

                    time.sleep(5)
            event.set_results({"output": stdout, "stderr": stderr})
            self.state.registered = True
        else:
            event.fail("Unit is not leader")
            return

    def on_add_gw_action(self, event):
        """Self-register for the AGW"""
        if self.unit.is_leader():
            agw_id = event.params["agw_id"]
            agw_name = event.params["agw_name"]
            orch_ip = event.params["orch_ip"]
            orch_net = event.params["orch_net"]
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run(
                "/home/magma/addgw.py --agw_id {} --agw_name {} --orch_ip {} --orch_net {}".format(
                    agw_id, agw_name, orch_ip, orch_net
                )
            )
            event.set_results({"output": stdout, "stderr": stderr})
        else:
            event.fail("Unit is not leader")
            return

    def on_reset_id_action(self, event):
        """Resets the hardware ID"""
        if self.unit.is_leader():
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run("sudo snowflake --force-new-key")
            event.set_results({"output": stdout, "stderr": stderr})
        else:
            event.fail("Unit is not leader")
            return

    def on_add_hosts_action(self, event):
        """Add Orchestrator host in /etc/hosts"""
        if self.unit.is_leader():
            orch_ip = event.params["orch_ip"]
            orch_hosts = "ORCH_IP controller.magma.test\nORCH_IP bootstrapper-controller.magma.test\nORCH_IP state-controller.magma.test\nORCH_IP dispatcher-controller.magma.test\nORCH_IP logger-controller.magma.test\nORCH_IP streamer-controller.magma.test\n"
            orch_hosts = orch_hosts.replace("ORCH_IP", orch_ip)
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run(
                "echo -e {} | sudo tee -a /etc/hosts".format(orch_hosts)
            )
            event.set_results({"output": stdout, "stderr": stderr})
        else:
            event.fail("Unit is not leader")
            return

    def on_restart_magma_action(self, event):
        """Resets the hardware ID"""
        if self.unit.is_leader():
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run("sudo service magma@* restart")
            event.set_results({"output": stdout, "stderr": stderr})
            if self.state.registered:
                self.state.ready = True
                self.send_relation_data()
        else:
            event.fail("Unit is not leader")
            return

    def on_del_gw_action(self, event):
        """Deregister from AGW"""
        if self.unit.is_leader():
            agw_id = event.params["agw_id"]
            orch_ip = event.params["orch_ip"]
            orch_net = event.params["orch_net"]
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run(
                "/home/magma/delgw.py --agw_id {} --orch_ip {} --orch_net {}".format(
                    agw_id, orch_ip, orch_net
                )
            )
            event.set_results({"output": stdout, "stderr": stderr})
        else:
            event.fail("Unit is not leader")
            return

    def on_add_test_subscriber_action(self, event):
        """Adds test subscriber to Orc8r HSS"""
        if self.unit.is_leader():
            orch_ip = event.params["orch_ip"]
            orch_net = event.params["orch_net"]
            proxy = self.get_ssh_proxy()
            stdout, stderr = proxy.run(
                "/home/magma/addtestsub.py --orch_ip {} --orch_net {}".format(
                    orch_ip, orch_net
                )
            )
            event.set_results({"output": stdout, "stderr": stderr})
        else:
            event.fail("Unit is not leader")
            return


if __name__ == "__main__":
    main(MagmaAGWProxyCharm)