Skip to content

Commit

Permalink
feat: Get sensible defaults for args
Browse files Browse the repository at this point in the history
  • Loading branch information
dorukozturk committed Dec 10, 2022
1 parent c74f80a commit d0e8fe5
Show file tree
Hide file tree
Showing 2 changed files with 63 additions and 18 deletions.
46 changes: 35 additions & 11 deletions hardeneks/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
from pkg_resources import resource_filename
import yaml

from botocore.exceptions import EndpointConnectionError
import boto3
import kubernetes
from rich.console import Console
import typer
Expand Down Expand Up @@ -34,22 +36,43 @@ def _config_callback(value: str):
return value


def _get_current_context(context):
if context:
return context
_, active_context = kubernetes.config.list_kube_config_contexts()
return active_context["name"]


def _get_namespaces(ignored_ns: list) -> list:
v1 = kubernetes.client.CoreV1Api()
namespaces = [i.metadata.name for i in v1.list_namespace().items]
return list(set(namespaces) - set(ignored_ns))


def _get_cluster_name(context, region):
try:
client = boto3.client("eks", region_name=region)
for name in client.list_clusters()["clusters"]:
if name in context:
return name
except EndpointConnectionError:
raise ValueError(f"{region} seems like a bad region name")


def _get_region():
return boto3.session.Session().region_name


@app.command()
def run_hardeneks(
region: str = typer.Option(
..., help="AWS region of the cluster. Ex: us-east-1"
default=None, help="AWS region of the cluster. Ex: us-east-1"
),
context: str = typer.Option(
...,
default=None,
help="K8s context.",
),
cluster: str = typer.Option(..., help="Cluster name."),
cluster: str = typer.Option(default=None, help="Cluster name."),
namespace: str = typer.Option(
default=None,
help="Specific namespace to harden. Default is all namespaces.",
Expand All @@ -59,9 +82,6 @@ def run_hardeneks(
callback=_config_callback,
help="Path to a hardeneks config file.",
),
kube_config: str = typer.Option(
default=None, help="Path to the kube config file."
),
):
"""
Main entry point to hardeneks.
Expand All @@ -72,24 +92,28 @@ def run_hardeneks(
cluster (str): Cluster name
namespace (str): Specific namespace to be checked
config (str): Path to hardeneks config file
kube_config (str): Path to the kube config file
Returns:
None
"""

kubernetes.config.load_kube_config(context=context)
context = _get_current_context(context)
if not cluster:
cluster = _get_cluster_name(context, region)

if not region:
region = _get_region()

console = Console()
console.rule("[b]HARDENEKS", characters="* ")
console.print(f"You are operating at {region}")
console.print(f"You context is {context}")
console.print(f"Your cluster name is {cluster}")
console.print(f"You are using {config} as your config file")
console.print()

kubernetes.config.load_kube_config(
config_file=kube_config, context=context
)

with open(config, "r") as f:
config = yaml.load(f, Loader=yaml.FullLoader)

Expand Down
35 changes: 28 additions & 7 deletions tests/test_cli.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,15 @@
from unittest.mock import patch
from pathlib import Path

from click import exceptions
import pytest
from typer.testing import CliRunner


from hardeneks import app, _config_callback

runner = CliRunner()
from hardeneks import (
_config_callback,
_get_cluster_name,
_get_current_context,
)


def test_config_callback_path_non_existent():
Expand Down Expand Up @@ -36,6 +38,25 @@ def test_config_callback_bad_yaml(tmp_path):
_config_callback(config)


def test_app_no_input():
result = runner.invoke(app, [])
assert result.exit_code == 2
@patch("kubernetes.config.list_kube_config_contexts")
def test_get_current_context_None(config):
config.return_value = ({}, {"name": "some-context"})
context = _get_current_context("")
assert context == "some-context"


def test_get_current_context():
context = "some-context"
assert _get_current_context(context) == context


@patch("boto3.client")
def test_get_cluster_name(client):
client.return_value.list_clusters.return_value = {
"clusters": ["gpu-cluster-test", "foo-cluster", "bad-cluster"]
}
context = "[email protected]"
region = "us-west-2"
cluster_name = "gpu-cluster-test"

assert _get_cluster_name(context, region) == cluster_name

0 comments on commit d0e8fe5

Please sign in to comment.