Skip to content

Commit

Permalink
feat: add first tests
Browse files Browse the repository at this point in the history
  • Loading branch information
Andreas Gruhler committed Nov 23, 2024
1 parent 327cf5e commit 17caba9
Show file tree
Hide file tree
Showing 8 changed files with 272 additions and 73 deletions.
6 changes: 6 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
**.swp
**.coverage
**.env
**.venv
**__pycache__
**vault_data
22 changes: 22 additions & 0 deletions kubernetes/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -51,3 +51,25 @@ mc undo my-snapshots/vault-snapshots-2f848f/vault_2024-09-06-1739.snapshot
mc ls --versions my-snapshots/vault-snapshots-2f848f
[2024-09-06 19:39:49 CEST] 28KiB Standard 1031052557042383613 v1 PUT vault_2024-09-06-1739.snapshot
```

## Development and tests

Requirements for running the mock server (`vault_server_mock.py`):
* HashiCorp Vault or OpenBao (`vault` binary)

To prepare the environment:
```bash
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```

Run the tests w/o coverage:
```bash
pytest
```

Run the tests with coverage:
```bash
coverage run .venv/bin/pytest
```
5 changes: 4 additions & 1 deletion kubernetes/requirements.txt
Original file line number Diff line number Diff line change
@@ -1,2 +1,5 @@
boto3
hvac
boto3
moto[s3]
pytest
coverage
72 changes: 0 additions & 72 deletions kubernetes/vault-snapshot.py

This file was deleted.

5 changes: 5 additions & 0 deletions kubernetes/vault_config.hcl
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
storage "raft" {
path = "./vault_data"
node_id = "devnode"
}
cluster_addr = "http://127.0.0.1:8201"
43 changes: 43 additions & 0 deletions kubernetes/vault_server_mock.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,43 @@
import subprocess

VAULT_PORT = 8200
VAULT_CONFIG = "./vault_config.hcl"
VAULT_TOKEN = "root"
VAULT_DATA_DIR = "./vault_data"

class VaultServer():
def reset_data(self, dir: str = VAULT_DATA_DIR):
"""
Reset Vault server mock raft data directory.
"""

subprocess.run(f"rm -rf {dir}/*", shell=True)
print(f"Vault data dir reset: {dir}")

def run(self, port: int = VAULT_PORT, config: str = VAULT_CONFIG,
token: str = VAULT_TOKEN):
"""
Start the Vault server mock with data dir and config.
"""

command = f"$(which vault) server -dev -dev-root-token-id={token} -config={config}"
self.proc = subprocess.Popen(command, shell=True)

def status(self):
"""
Return Vault server mock status.
A None value indicates that the process hadn’t yet terminated at the
time of the last method call:
* https://docs.python.org/3/library/subprocess.html#subprocess.Popen.returncode
"""
return self.proc.returncode

def stop(self):
"""
Kill Vault server mock process.
"""

self.proc.kill()
self.proc.wait()
print(f"Process returned with return code: {self.proc.returncode}")
132 changes: 132 additions & 0 deletions kubernetes/vault_snapshot.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
#!/usr/bin/env python

import logging
import boto3
from botocore.exceptions import ClientError
from hvac.api.auth_methods import Kubernetes
import hvac
import os
import datetime

class VaultSnapshot:
"""
Create Vault snapshots on S3.
"""

def __init__(self, **kwargs):
"""
Init S3 and hvac clients
"""

# setup logger
self.logger = logging.getLogger(__name__)

# read input keyword arguments
if "vault_addr" in kwargs:
self.vault_addr = kwargs['vault_addr']
elif "VAULT_ADDR" in os.environ:
self.vault_addr = os.environ['VAULT_ADDR']
else:
raise NameError("VAULT_ADDR undefined")

if "vault_token" in kwargs:
self.vault_token = kwargs['vault_token']
elif "VAULT_TOKEN" in os.environ:
self.vault_token = os.environ['VAULT_TOKEN']
else:
raise NameError("VAULT_TOKEN undefined")

if "s3_access_key_id" in kwargs:
self.s3_access_key_id = kwargs['s3_access_key_id']
elif "AWS_ACCESS_KEY_ID" in os.environ:
self.s3_access_key_id = os.environ['AWS_ACCESS_KEY_ID']
else:
raise NameError("AWS_ACCESS_KEY_ID undefined")

if "s3_secret_access_key" in kwargs:
self.s3_secret_access_key = kwargs['s3_secret_access_key']
elif "AWS_SECRET_ACCESS_KEY" in os.environ:
self.s3_secret_access_key = os.environ['AWS_SECRET_ACCESS_KEY']
else:
raise NameError("AWS_SECRET_ACCESS_KEY undefined")

if "s3_host" in kwargs:
self.s3_host = kwargs['s3_host']
elif "S3_HOST" in os.environ:
self.s3_host = os.environ['S3_HOST']
else:
raise NameError("S3_HOST undefined")

if "s3_bucket" in kwargs:
self.s3_bucket = kwargs['s3_bucket']
elif "S3_BUCKET" in os.environ:
self.s3_bucket = os.environ['S3_BUCKET']
else:
raise NameError("S3_BUCKET undefined")

# Boto S3 client
# * https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html
self.s3_client = boto3.client(service_name='s3',
endpoint_url=self.s3_host,
aws_access_key_id=self.s3_access_key_id,
aws_secret_access_key=self.s3_secret_access_key)

# Authenticate with Kubernetes ServiceAccount if vault_token is empty
# https://hvac.readthedocs.io/en/stable/usage/auth_methods/kubernetes.html
#hvac_client = hvac.Client(url=url, verify=certificate_path)
#f = open('/var/run/secrets/kubernetes.io/serviceaccount/token')
#jwt = f.read()
####VAULT_TOKEN=$(vault write -field=token auth/kubernetes/login role="${VAULT_ROLE}" jwt="${JWT}")
#Kubernetes(hvac_client.adapter).login(role=role, jwt=jwt)

self.logger.debug(f"Connecting to Vault API {self.vault_addr}")
self.hvac_client = hvac.Client(url=self.vault_addr)
self.hvac_client.token = self.vault_token

assert self.hvac_client.is_authenticated()

def snapshot(self):
"""Create Vault integrated storage (Raft) snapshot.
The snapshot is returned as binary data and should be redirected to
a file:
* https://developer.hashicorp.com/vault/api-docs/system/storage/raft
* https://hvac.readthedocs.io/en/stable/source/hvac_api_system_backend.html
"""

with self.hvac_client.sys.take_raft_snapshot() as resp:
assert resp.ok

self.logger.debug("Raft snapshot status code: %d" % resp.status_code)

date_str = datetime.datetime.now(datetime.UTC).strftime("%F-%H%M")
file_name = "vault_%s.snapshot" % (date_str)
self.logger.debug(f"File name: {file_name}")

# Upload the file
# * https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/s3/client/put_object.html
try:
response = self.s3_client.put_object(
Body=resp.content,
Bucket=self.s3_bucket,
Key=file_name,
)
self.logger.debug("s3 put_object response: %s", response)
except ClientError as e:
logging.error(e)

# Iterate and remove expired snapshots:
# https://boto3.amazonaws.com/v1/documentation/api/latest/guide/migrations3.html
s3 = boto3.resource(service_name='s3',
endpoint_url=self.s3_host,
aws_access_key_id=self.s3_access_key_id,
aws_secret_access_key=self.s3_secret_access_key)
bucket = s3.Bucket(self.s3_bucket)
for key in bucket.objects.all():
self.logger.debug(key.key)
# todo: do the S3_EXPIRE_DAYS magic

return file_name

if __name__=="__main__":
VaultSnapshot.snapshot()
60 changes: 60 additions & 0 deletions kubernetes/vault_snapshot_test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
import pytest
import boto3
import hvac
from moto import mock_aws
from unittest.mock import patch, create_autospec

from vault_snapshot import VaultSnapshot
from vault_server_mock import VaultServer

class TestVaultSnapshots:
"""
Test Vault snapshot functionality.
"""

@pytest.fixture
def mock_vault_server(self):
"""
Run a Vault server mock.
"""
self.mock = VaultServer()
self.mock.reset_data()
# run mock
self.mock.run()
# return process status, 'None' means process is still running
yield self.mock.status()
# when tests are done, teardown the Vault server
self.mock.stop()

@mock_aws
def test_snapshot(self, mock_vault_server):
"""
Test snapshot functionality using boto3 and moto.
https://docs.getmoto.org/en/latest/docs/getting_started.html#decorator
"""

print(f"The current Vault server mock process status is: {mock_vault_server}")

bucket_name = "vault-snapshots"
region_name = "us-east-1"
conn = boto3.resource("s3", region_name=region_name)
# We need to create the bucket since this is all in Moto's 'virtual' AWS account
conn.create_bucket(Bucket=bucket_name)

vault_snapshot = VaultSnapshot(
vault_addr="http://127.0.0.1:8200",
vault_token="root",
s3_access_key_id="test",
s3_secret_access_key="test",
s3_host=f"https://s3.{region_name}.amazonaws.com",
s3_bucket=bucket_name
)
file_name = vault_snapshot.snapshot()

body = conn.Object(bucket_name,
file_name).get()#["Body"].read()#.decode("utf-8")

#print(body)

#assert body == "is awesome"

0 comments on commit 17caba9

Please sign in to comment.