-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Andreas Gruhler
committed
Nov 23, 2024
1 parent
327cf5e
commit 17caba9
Showing
8 changed files
with
272 additions
and
73 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,6 @@ | ||
**.swp | ||
**.coverage | ||
**.env | ||
**.venv | ||
**__pycache__ | ||
**vault_data |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,5 @@ | ||
boto3 | ||
hvac | ||
boto3 | ||
moto[s3] | ||
pytest | ||
coverage |
This file was deleted.
Oops, something went wrong.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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}") |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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" |