├── .gitignore ├── README.md ├── common.py ├── locustfile.py ├── locusts ├── __init__.py ├── auth_approle.py ├── auth_userpass.py ├── dyn_mongodb.py ├── dyn_mysql.py ├── key_value.py ├── pki.py ├── totp.py └── transit.py ├── prepare.py └── requirements.txt /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | env/ 12 | build/ 13 | develop-eggs/ 14 | dist/ 15 | downloads/ 16 | eggs/ 17 | .eggs/ 18 | lib/ 19 | lib64/ 20 | parts/ 21 | sdist/ 22 | var/ 23 | wheels/ 24 | *.egg-info/ 25 | .installed.cfg 26 | *.egg 27 | 28 | # PyInstaller 29 | # Usually these files are written by a python script from a template 30 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 31 | *.manifest 32 | *.spec 33 | 34 | # Installer logs 35 | pip-log.txt 36 | pip-delete-this-directory.txt 37 | 38 | # Unit test / coverage reports 39 | htmlcov/ 40 | .tox/ 41 | .coverage 42 | .coverage.* 43 | .cache 44 | nosetests.xml 45 | coverage.xml 46 | *.cover 47 | .hypothesis/ 48 | 49 | # Translations 50 | *.mo 51 | *.pot 52 | 53 | # Django stuff: 54 | *.log 55 | local_settings.py 56 | 57 | # Flask stuff: 58 | instance/ 59 | .webassets-cache 60 | 61 | # Scrapy stuff: 62 | .scrapy 63 | 64 | # Sphinx documentation 65 | docs/_build/ 66 | 67 | # PyBuilder 68 | target/ 69 | 70 | # Jupyter Notebook 71 | .ipynb_checkpoints 72 | 73 | # pyenv 74 | .python-version 75 | 76 | # celery beat schedule file 77 | celerybeat-schedule 78 | 79 | # SageMath parsed files 80 | *.sage.py 81 | 82 | # dotenv 83 | .env 84 | 85 | # virtualenv 86 | .venv 87 | venv/ 88 | ENV/ 89 | 90 | # Spyder project settings 91 | .spyderproject 92 | .spyproject 93 | 94 | # Rope project settings 95 | .ropeproject 96 | 97 | # mkdocs documentation 98 | /site 99 | 100 | # mypy 101 | .mypy_cache/ 102 | 103 | .idea/ 104 | testdata.json 105 | 106 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | Load Tests for Vault 2 | ==================== 3 | 4 | This project aims to generate realistic load against Vault (and, by extension, 5 | Consul) by exercising various secrets engines and auth methods. 6 | 7 | ## Prerequisites 8 | 9 | You need to have Vault running, obviously, and the vault must be unsealed before 10 | starting the test. 11 | 12 | The code in this repo requires Python3.6. 13 | 14 | If you want to test dynamic secret generation via the database backends (MySQL, 15 | MongoDB), you must have those database engines running during the test. You will 16 | need to set up environment variables with the connection strings in the appropriate 17 | formats, for example: 18 | 19 | export MONGODB_URL="mongodb://localhost:27017/admin" 20 | export MYSQL_URL="root:password@tcp(127.0.0.1:3306)/mysql" 21 | 22 | If you don't have the databases available, remove the database locusts from the 23 | `locustfile`. 24 | 25 | ## Setup 26 | 27 | 1. Clone this repo. 28 | 2. `pip install -r requirements.txt`. 29 | 3. Run the `prepare.py` script to populate Vault with random secrets. 30 | 4. Run `locust` to start a test. 31 | 32 | 33 | ## Preparing Vault for the test 34 | 35 | The `prepare.py` script fills Vault with a bunch of random secrets that are 36 | then queried during the test. The paths are random strings of hex characters 37 | with two levels, like this: 38 | 39 | fd/410a7adef1cd7ce6fbfaebc3b4eb49b9 40 | e3/c0a0bafa86e81e2b4b8775f9b9dacfcf 41 | fd/8f4293e2fd64cadda5f6cb936ff054e4 42 | 43 | Each secret consists of a single key whose value is a string of random bytes. 44 | The length of the string is controlled by the `--secret-size` parameter 45 | (default: 1024 bytes), and the number of secrets created is controlled by the 46 | `--num-secrets` parameter (default: 1000). 47 | 48 | The script assumes that the target Vault instance lives at 49 | http://localhost:8200. If this is not correct, use the `--host` or `-h` option 50 | to specify the correct URL. 51 | 52 | The final argument to `prepare.py` must be a Vault token for authentication. 53 | The token can instead be passed in the environment variable `VAULT_TOKEN` if you 54 | prefer. 55 | 56 | Example usage: 57 | 58 | VAULT_TOKEN="72e7ff6e-8f44-5414-75a8-99d308649954" ./prepare.py --num-secrets=1000000 --host="http://localhost:8200" 59 | 60 | 61 | ## Running the test 62 | 63 | Run the `locust` command to start an interactive web interface at 64 | http://localhost:8089. If you needed to use the `--host` parameter above, 65 | make sure you use it here too. The test will not start until you open the 66 | web page and click the **Start Swarming** button. 67 | 68 | If you'd prefer to run the test non-interactively, you can specify additional 69 | command-line parameters along with `--no-web`. 70 | 71 | For example, to start a test with 25 locusts (25 simulated users), starting 5 72 | per second until all 25 are running: 73 | 74 | locust -H http://localhost:8200 -c 25 -r 5 --no-web 75 | 76 | ### testdata.json 77 | 78 | Running the ./prepare.py script creates a testdata.json file on the host it is run from. 79 | 80 | This file contains the vault token used to run the pre-fill vault with secrets, so treat that file as secret. 81 | 82 | If you plan to run the load test from multiple hosts, you will need to copy the testdata.json file to those hosts 83 | before running the load test in order for it to work correctly. 84 | 85 | 86 | ## What the test does 87 | 88 | The test simulates several different patterns of access: 89 | 90 | 1. Reading, writing, and listing secrets with the KV engine. 91 | 2. Generating certificates from the PKI engine. 92 | 3. Encrypting data via the Transit engine. 93 | 4. Generating dynamic secrets from the MySQL, MongoDB, and TOTP engines. 94 | 5. Authenticating via username/password and AppRole. 95 | 96 | The tests are weighted so that 60% of the users are interacting with the KV 97 | engine, 20% with the PKI engine, and 20% with the Transit engine. These weights 98 | can be adjusted by changing the `weight` parameter in each of the files in 99 | the `locusts` folder. 100 | 101 | ## Future enhancements 102 | 103 | - [x] Authentication 104 | - [x] Dynamic secret generation 105 | - [ ] Reports and analysis? 106 | 107 | -------------------------------------------------------------------------------- /common.py: -------------------------------------------------------------------------------- 1 | import random 2 | import string 3 | from locust.clients import HttpSession 4 | import requests 5 | 6 | 7 | def random_hex(size) -> str: 8 | return ''.join([random.choice(string.hexdigits) for _ in range(size)]).lower() 9 | 10 | 11 | def key_path_1() -> str: 12 | return random_hex(2) 13 | 14 | 15 | def key_path_2() -> str: 16 | return random_hex(32) 17 | 18 | 19 | def key_path() -> str: 20 | return '%s/%s' % (key_path_1(), key_path_2()) 21 | 22 | 23 | def random_data(size) -> str: 24 | return ''.join([random.choice(string.ascii_letters) for _ in range(size)]) 25 | 26 | 27 | def get_kv_version(client: requests.Session=None, host: str=None, token: str=None) -> int: 28 | if isinstance(client, HttpSession): 29 | s = client 30 | r = s.get('/v1/sys/mounts') 31 | elif isinstance(client, requests.Session): 32 | s = client 33 | r = s.get(f'{host}/v1/sys/mounts') 34 | else: 35 | s = requests.Session() 36 | s.headers = {'X-Vault-Token': token} 37 | r = s.get(f'{host}/v1/sys/mounts') 38 | 39 | r.raise_for_status() 40 | 41 | version = 1 42 | for key, val in r.json().items(): 43 | if key == 'secret/': 44 | if 'options' in val: 45 | version = int(val['options'].get('version', 1)) 46 | break 47 | 48 | return version 49 | -------------------------------------------------------------------------------- /locustfile.py: -------------------------------------------------------------------------------- 1 | from locusts.key_value import KeyValueLocust 2 | from locusts.transit import TransitLocust 3 | from locusts.pki import PkiLocust 4 | from locusts.dyn_mongodb import MongoDbLocust 5 | from locusts.dyn_mysql import MysqlLocust 6 | from locusts.totp import TotpLocust 7 | from locusts.auth_userpass import UserPassAuthLocust 8 | from locusts.auth_approle import AppRoleLocust 9 | 10 | __static__ = [KeyValueLocust, TransitLocust, PkiLocust] 11 | __dynamic__ = [MysqlLocust, MongoDbLocust, TotpLocust] 12 | __auth__ = [UserPassAuthLocust, AppRoleLocust] 13 | 14 | __all__ = __static__ + __dynamic__ + __auth__ 15 | 16 | # import logging 17 | # logging.getLogger().setLevel(logging.DEBUG) 18 | # requests_log = logging.getLogger("requests.packages.urllib3") 19 | # requests_log.setLevel(logging.DEBUG) 20 | # requests_log.propagate = True 21 | # 22 | # import http.client 23 | # http.client.HTTPConnection.debuglevel = 1 24 | -------------------------------------------------------------------------------- /locusts/__init__.py: -------------------------------------------------------------------------------- 1 | from locust import TaskSet, HttpLocust 2 | from locust.clients import ResponseContextManager, HttpSession, RequestException, events, CatchResponseError 3 | import time 4 | 5 | import os 6 | import json 7 | 8 | 9 | class VaultLocust(HttpLocust): 10 | 11 | token = None 12 | testdata = None 13 | 14 | def __init__(self): 15 | super().__init__() 16 | self.client = VaultSession(base_url=self.host) 17 | 18 | def setup(self): 19 | with open('testdata.json', 'r') as f: 20 | config = json.load(f) 21 | self._set_testdata(config) 22 | self._set_token(config['token']) 23 | 24 | @classmethod 25 | def _set_token(cls, token): 26 | cls.token = token 27 | 28 | @classmethod 29 | def _set_testdata(cls, data): 30 | cls.testdata = data 31 | 32 | 33 | class VaultTaskSet(TaskSet): 34 | 35 | def mount(self, name: str, mount_point: str=None): 36 | mount_point = mount_point or name 37 | r = self.client.get('/v1/sys/mounts') 38 | if f'{mount_point}/' not in r.json(): 39 | self.client.post(f'/v1/sys/mounts/{mount_point}', json={'type': name}) 40 | 41 | def enable_auth(self, name: str, path: str=None): 42 | path = path or name 43 | r = self.client.get('/v1/sys/auth') 44 | if f'{path}/' not in r.json(): 45 | self.client.post(f'/v1/sys/auth/{path}', json={'type': name}) 46 | 47 | def revoke_lease(self, lease_id: str): 48 | self.client.put('/v1/sys/leases/revoke', 49 | json={'lease_id': lease_id}) 50 | 51 | def is_in_list(self, key: str, uri: str) -> bool: 52 | with self.client.request('LIST', uri, catch_response=True) as r: 53 | if r.status_code == 404: 54 | r.success() 55 | return False 56 | else: 57 | return key in r.json()['data']['keys'] 58 | 59 | @property 60 | def client(self) -> HttpSession: 61 | client = super().client # type: HttpSession 62 | client.headers['X-Vault-Token'] = self.locust.token 63 | return client 64 | 65 | 66 | class VaultSession(HttpSession): 67 | 68 | def request(self, method, url, name=None, catch_response=False, **kwargs): 69 | 70 | # Load any TLS certificates specified in the VAULT_CACERT env var 71 | self.verify = os.environ.get('VAULT_CACERT', None) 72 | 73 | # prepend url with hostname unless it's already an absolute URL 74 | url = self._build_url(url) 75 | 76 | # store meta data that is used when reporting the request to locust's statistics 77 | request_meta = dict() 78 | 79 | # set up pre_request hook for attaching meta data to the request object 80 | request_meta["method"] = method 81 | request_meta["start_time"] = time.time() 82 | 83 | response = self._send_request_safe_mode(method, url, **kwargs) 84 | 85 | # record the consumed time 86 | request_meta["response_time"] = int((time.time() - request_meta["start_time"]) * 1000) 87 | 88 | request_meta["name"] = name or (response.history and response.history[0] or response).request.path_url 89 | 90 | # get the length of the content, but if the argument stream is set to True, we take 91 | # the size from the content-length header, in order to not trigger fetching of the body 92 | if kwargs.get("stream", False): 93 | request_meta["content_size"] = int(response.headers.get("content-length") or 0) 94 | else: 95 | request_meta["content_size"] = len(response.content or "") 96 | 97 | if catch_response: 98 | response.locust_request_meta = request_meta 99 | return ResponseContextManager(response) 100 | else: 101 | try: 102 | response.raise_for_status() 103 | except RequestException as e: 104 | try: 105 | e = CatchResponseError('. '.join(response.json()['errors'])) 106 | except KeyError: 107 | e = CatchResponseError(e) 108 | except json.JSONDecodeError: 109 | e = CatchResponseError(e) 110 | 111 | events.request_failure.fire( 112 | request_type=request_meta["method"], 113 | name=request_meta["name"], 114 | response_time=request_meta["response_time"], 115 | exception=e, 116 | ) 117 | else: 118 | events.request_success.fire( 119 | request_type=request_meta["method"], 120 | name=request_meta["name"], 121 | response_time=request_meta["response_time"], 122 | response_length=request_meta["content_size"], 123 | ) 124 | return response 125 | -------------------------------------------------------------------------------- /locusts/auth_approle.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import random 4 | 5 | from locust import task 6 | 7 | sys.path.append(os.getcwd()) 8 | from locusts import VaultTaskSet, VaultLocust 9 | 10 | 11 | class AppRoleTaskSet(VaultTaskSet): 12 | ROLE_NAME = 'test-approle' 13 | SECRET_NAME = 'test-secret' 14 | role_id = None 15 | secrets = [] 16 | 17 | def setup(self): 18 | self.enable_auth('approle') 19 | self.create_approle() 20 | 21 | def teardown(self): 22 | self.client.delete(f'/v1/auth/approle/role/{self.ROLE_NAME}') 23 | 24 | def create_approle(self): 25 | self.client.post(f'/v1/auth/approle/role/{self.ROLE_NAME}') 26 | r = self.client.get(f'/v1/auth/approle/role/{self.ROLE_NAME}/role-id') 27 | self._set_roleid(r.json()['data']['role_id']) 28 | 29 | @classmethod 30 | def _set_roleid(cls, role_id): 31 | cls.role_id = role_id 32 | 33 | @classmethod 34 | def _append_secret(cls, secret_id, accessor): 35 | cls.secrets.append((secret_id, accessor)) 36 | 37 | @task 38 | def create_secret(self): 39 | r = self.client.post(f'/v1/auth/approle/role/{self.ROLE_NAME}/secret-id') 40 | secret_id = r.json()['data']['secret_id'] 41 | accessor = r.json()['data']['secret_id_accessor'] 42 | self._append_secret(secret_id, accessor) 43 | 44 | @task 45 | def auth_success(self): 46 | try: 47 | secret = random.choice(self.secrets) 48 | self.client.post('/v1/auth/approle/login', 49 | json={'role_id': self.role_id, 50 | 'secret_id': secret[0]}) 51 | except IndexError: 52 | pass 53 | 54 | @task 55 | def auth_failure(self): 56 | try: 57 | secret = random.choice(self.secrets) 58 | with self.client.post('/v1/auth/approle/login', 59 | json={'role_id': self.role_id, 60 | 'secret_id': secret[0] + 'XXX'}, 61 | catch_response=True) as r: 62 | if r.status_code == 400 and ('failed to validate credentials' in r.json()['errors'][0] 63 | or 'invalid secret id' in r.json()['errors'][0]): 64 | r.success() 65 | else: 66 | r.failure('unexpected response to bad auth token: ' + r.content.decode('utf-8')) 67 | except IndexError: 68 | pass 69 | 70 | 71 | class AppRoleLocust(VaultLocust): 72 | task_set = AppRoleTaskSet 73 | weight = 1 74 | min_wait = 5000 75 | max_wait = 10000 76 | -------------------------------------------------------------------------------- /locusts/auth_userpass.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from locust import task 5 | 6 | sys.path.append(os.getcwd()) 7 | from common import random_data 8 | from locusts import VaultTaskSet, VaultLocust 9 | 10 | 11 | class UserPassAuthTasks(VaultTaskSet): 12 | USER_NAME = 'test-user' 13 | password = None 14 | 15 | def setup(self): 16 | self.enable_auth('userpass') 17 | self.reset_password() 18 | 19 | def teardown(self): 20 | self.client.delete(f'/v1/auth/userpass/users/{self.USER_NAME}') 21 | 22 | @classmethod 23 | def _set_password(cls): 24 | cls.password = random_data(16) 25 | return cls.password 26 | 27 | @task 28 | def auth_success(self): 29 | self.client.post(f'/v1/auth/userpass/login/{self.USER_NAME}', 30 | json={'password': UserPassAuthTasks.password}) 31 | 32 | @task 33 | def auth_failure(self): 34 | with self.client.post(f'/v1/auth/userpass/login/{self.USER_NAME}', 35 | json={'password': random_data(16)}, 36 | catch_response=True) as r: 37 | if r.status_code == 400 and 'invalid username or password' in r.json()['errors']: 38 | r.success() 39 | else: 40 | r.failure('unexpected response to invalid login') 41 | 42 | @task 43 | def reset_password(self): 44 | self.client.post(f'/v1/auth/userpass/users/{self.USER_NAME}', 45 | json={'password': UserPassAuthTasks._set_password()}) 46 | 47 | 48 | class UserPassAuthLocust(VaultLocust): 49 | task_set = UserPassAuthTasks 50 | weight = 1 51 | min_wait = 5000 52 | max_wait = 10000 53 | -------------------------------------------------------------------------------- /locusts/dyn_mongodb.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from locust import task 5 | 6 | sys.path.append(os.getcwd()) 7 | from locusts import VaultTaskSet, VaultLocust 8 | 9 | 10 | class MongoDbTasks(VaultTaskSet): 11 | """ 12 | In order for this test to work, you must have MongoDB running somewhere that is accessible to Vault. 13 | Set the environment variable MONGODB_URL to point to the MongoDB instance. If you don't have MongoDB, remove 14 | MongoDbLocust from the locustfile. 15 | """ 16 | CONFIG_NAME = 'test-mongodb-local' 17 | ROLE_NAME = 'test-mongodb-role' 18 | DEFAULT_CONN_URL = 'mongodb://localhost:27017/admin' 19 | CREATE_SQL = '{"db": "admin", "roles": [{"role": "read", "db": "foo"}]}' 20 | 21 | def setup(self): 22 | self._set_conn_url(os.environ.get('MONGODB_URL', self.DEFAULT_CONN_URL)) 23 | self.mount('database') 24 | self.create_connection() 25 | self.create_role() 26 | 27 | def teardown(self): 28 | self.delete_role() 29 | self.delete_connection() 30 | 31 | @classmethod 32 | def _set_conn_url(cls, url): 33 | cls.conn_url = url 34 | 35 | def create_connection(self): 36 | if self.is_in_list(self.CONFIG_NAME, '/v1/database/config'): 37 | self.delete_connection() 38 | self.client.post(f'/v1/database/config/{self.CONFIG_NAME}', 39 | json={'plugin_name': 'mongodb-database-plugin', 40 | 'allowed_roles': self.ROLE_NAME, 41 | 'connection_url': self.conn_url}) 42 | 43 | def create_role(self): 44 | if self.is_in_list(self.ROLE_NAME, '/v1/database/roles'): 45 | self.delete_role() 46 | self.client.post(f'/v1/database/roles/{self.ROLE_NAME}', 47 | json={'db_name': self.CONFIG_NAME, 48 | 'creation_statements': self.CREATE_SQL, 49 | 'default_ttl': '1h', 50 | 'max_ttl': '72h'}) 51 | 52 | def delete_connection(self): 53 | self.client.delete(f'/v1/database/config/{self.CONFIG_NAME}') 54 | 55 | def delete_role(self): 56 | self.client.delete(f'/v1/database/roles/{self.ROLE_NAME}') 57 | 58 | @task 59 | def generate_creds(self): 60 | r = self.client.get(f'/v1/database/creds/{self.ROLE_NAME}') 61 | lease_id = r.json()['lease_id'] 62 | self.client.put('/v1/sys/leases/revoke', 63 | json={'lease_id': lease_id}) 64 | 65 | 66 | class MongoDbLocust(VaultLocust): 67 | task_set = MongoDbTasks 68 | weight = 1 69 | min_wait = 5000 70 | max_wait = 10000 71 | -------------------------------------------------------------------------------- /locusts/dyn_mysql.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | import locust 5 | from locust import task 6 | 7 | sys.path.append(os.getcwd()) 8 | from locusts import VaultTaskSet, VaultLocust 9 | 10 | 11 | class MysqlTasks(VaultTaskSet): 12 | """ 13 | In order for this test to work, you must have MySQL running somewhere that is accessible to Vault. 14 | Set the environment variable MYSQL_URL to point to the MySQL instance. If you don't have MySQL, remove 15 | MysqlLocust from the locustfile. 16 | """ 17 | CONFIG_NAME = 'test-mysql-local' 18 | ROLE_NAME = 'test-mysql-role' 19 | DEFAULT_CONN_URL = '/mysql' 20 | CREATE_SQL = "CREATE USER '{{name}}'@'%' IDENTIFIED BY '{{password}}'; GRANT SELECT ON *.* TO '{{name}}'@'%';" 21 | conn_url = None 22 | 23 | def setup(self): 24 | self._set_conn_url(os.environ.get('MYSQL_URL', self.DEFAULT_CONN_URL)) 25 | self.mount('database') 26 | self.create_connection() 27 | self.create_role() 28 | 29 | def teardown(self): 30 | self.delete_role() 31 | self.delete_connection() 32 | 33 | @classmethod 34 | def _set_conn_url(cls, url): 35 | cls.conn_url = url 36 | 37 | def create_connection(self): 38 | if self.is_in_list(self.CONFIG_NAME, '/v1/database/config'): 39 | self.delete_connection() 40 | self.client.post(f'/v1/database/config/{self.CONFIG_NAME}', 41 | json={'plugin_name': 'mysql-database-plugin', 42 | 'allowed_roles': self.ROLE_NAME, 43 | 'connection_url': self.conn_url}) 44 | 45 | def delete_connection(self): 46 | self.client.delete(f'/v1/database/config/{self.CONFIG_NAME}') 47 | 48 | def create_role(self): 49 | if self.is_in_list(self.ROLE_NAME, '/v1/database/roles'): 50 | self.delete_role() 51 | self.client.post(f'/v1/database/roles/{self.ROLE_NAME}', 52 | json={'db_name': self.CONFIG_NAME, 53 | 'creation_statements': self.CREATE_SQL, 54 | 'default_ttl': '1h', 55 | 'max_ttl': '72h'}) 56 | 57 | def delete_role(self): 58 | self.client.delete(f'/v1/database/roles/{self.ROLE_NAME}') 59 | 60 | @task 61 | def generate_creds(self): 62 | r = self.client.get(f'/v1/database/creds/{self.ROLE_NAME}') 63 | if 'lease_id' in r.json(): 64 | lease_id = r.json()['lease_id'] 65 | self.client.put('/v1/sys/leases/revoke', 66 | json={'lease_id': lease_id}) 67 | 68 | 69 | class MysqlLocust(VaultLocust): 70 | task_set = MysqlTasks 71 | weight = 1 72 | min_wait = 5000 73 | max_wait = 10000 74 | -------------------------------------------------------------------------------- /locusts/key_value.py: -------------------------------------------------------------------------------- 1 | import json 2 | import sys 3 | 4 | import os 5 | import random 6 | 7 | from locust import HttpLocust, task 8 | 9 | sys.path.append(os.getcwd()) 10 | import common 11 | 12 | from locusts import VaultTaskSet, VaultLocust 13 | 14 | 15 | class KeyValueTasks(VaultTaskSet): 16 | 17 | def __init__(self, parent): 18 | super().__init__(parent) 19 | self.kv_version = 1 20 | 21 | def on_start(self): 22 | self.kv_version = common.get_kv_version(client=self.client) 23 | 24 | @task 25 | def get_kv_secret(self): 26 | key = random.choice(self.locust.testdata['keys']) 27 | if self.kv_version == 1: 28 | self.client.get(f'/v1/secret/test/{key}', name='/v1/secret/[key1]/[key2]') 29 | else: 30 | self.client.get(f'/v1/secret/data/test/{key}', name='/v1/secret/[key1]/[key2]') 31 | 32 | @task 33 | def put_kv_secret(self): 34 | key = random.choice(self.locust.testdata['keys']) 35 | payload = common.random_data(self.locust.testdata['secret_size']) 36 | if self.kv_version == 1: 37 | self.client.put(f'/v1/secret/test/{key}', 38 | json={'value': payload}, 39 | name='/v1/secret/test/[key1]/[key2]') 40 | else: 41 | self.client.put(f'/v1/secret/data/test/{key}', 42 | json={'data': {'value': payload}}, 43 | name='/v1/secret/test/[key1]/[key2]') 44 | 45 | @task 46 | def list_l1_secrets(self): 47 | if self.kv_version == 1: 48 | self.client.request('LIST', '/v1/secret/test', name='/v1/secret/test') 49 | else: 50 | self.client.request('LIST', '/v1/secret/metadata/test', name='/v1/secret/test') 51 | 52 | @task 53 | def list_l2_secrets(self): 54 | key_path = common.key_path_1() 55 | if self.kv_version == 1: 56 | self.client.request('LIST', f'/v1/secret/test/{key_path}', name='/v1/secret/test/[key1]') 57 | else: 58 | self.client.request('LIST', f'/v1/secret/metadata/test/{key_path}', name='/v1/secret/test/[key1]') 59 | 60 | 61 | class KeyValueLocust(VaultLocust): 62 | task_set = KeyValueTasks 63 | weight = 3 64 | min_wait = 5000 65 | max_wait = 10000 66 | -------------------------------------------------------------------------------- /locusts/pki.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from locust import HttpLocust, task 5 | 6 | sys.path.append(os.getcwd()) 7 | from locusts import VaultTaskSet, VaultLocust 8 | 9 | 10 | class PkiTasks(VaultTaskSet): 11 | DOMAIN_NAME = 'example.com' 12 | ROLE_NAME = 'test-pki-role' 13 | 14 | def setup(self): 15 | self.mount('pki') 16 | self.client.post('/v1/pki/root/generate/internal', 17 | json={'common_name': self.DOMAIN_NAME, 'ttl': '8760h'}) 18 | self.create_role() 19 | 20 | def teardown(self): 21 | self.delete_role() 22 | 23 | def create_role(self): 24 | if self.is_in_list(self.ROLE_NAME, '/v1/pki/roles'): 25 | self.delete_role() 26 | self.client.post(f'/v1/pki/roles/{self.ROLE_NAME}', 27 | json={'allowed_domains': self.DOMAIN_NAME, 28 | 'max_ttl': '72h', 29 | 'allow_subdomains': True}) 30 | 31 | def delete_role(self): 32 | self.client.delete(f'/v1/pki/roles/{self.ROLE_NAME}') 33 | 34 | @task 35 | def generate_cert(self): 36 | self.client.post('/v1/pki/issue/test-pki-role', 37 | json={'common_name': f'foo.{self.DOMAIN_NAME}'}) 38 | 39 | 40 | class PkiLocust(VaultLocust): 41 | task_set = PkiTasks 42 | weight = 1 43 | min_wait = 5000 44 | max_wait = 10000 45 | -------------------------------------------------------------------------------- /locusts/totp.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | import common 4 | 5 | from locust import task 6 | from locusts import VaultTaskSet, VaultLocust 7 | 8 | 9 | class TotpTasks(VaultTaskSet): 10 | KEY_NAME = 'test-totp-key' 11 | ISSUER = 'Vault' 12 | ACCOUNT_NAME = 'test@test.com' 13 | 14 | def setup(self): 15 | self.mount('totp') 16 | if not self.is_in_list(self.KEY_NAME, '/v1/totp/keys'): 17 | self.client.post(f'/v1/totp/keys/{self.KEY_NAME}', 18 | json={'generate': True, 19 | 'issuer': self.ISSUER, 20 | 'account_name': self.ACCOUNT_NAME}) 21 | 22 | @task 23 | def generate(self): 24 | r = self.client.get(f'/v1/totp/code/{self.KEY_NAME}') 25 | 26 | 27 | class TotpLocust(VaultLocust): 28 | task_set = TotpTasks 29 | weight = 1 30 | min_wait = 5000 31 | max_wait = 60000 32 | -------------------------------------------------------------------------------- /locusts/transit.py: -------------------------------------------------------------------------------- 1 | import json 2 | import sys 3 | 4 | import base64 5 | import os 6 | 7 | from locust import HttpLocust, task 8 | from locust.clients import HttpSession 9 | 10 | sys.path.append(os.getcwd()) 11 | import common 12 | 13 | from locusts import VaultTaskSet, VaultLocust 14 | 15 | 16 | class TransitTasks(VaultTaskSet): 17 | 18 | def setup(self): 19 | self.mount('transit') 20 | 21 | @task 22 | def encrypt_block(self): 23 | data = common.random_data(self.locust.testdata['transit_size']) 24 | data = base64.b64encode(data.encode()).decode() 25 | self.client.post('/v1/transit/encrypt/test', json={'plaintext': data}) 26 | 27 | 28 | class TransitLocust(VaultLocust): 29 | task_set = TransitTasks 30 | weight = 1 31 | min_wait = 5000 32 | max_wait = 10000 33 | -------------------------------------------------------------------------------- /prepare.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import os 4 | import common 5 | import requests 6 | import click 7 | import json 8 | from typing import List 9 | from common import get_kv_version 10 | 11 | 12 | def populate(host: str, count: int, size: int, cacerts: str, token: str) -> List[str]: 13 | 14 | s = requests.Session() 15 | s.headers = {'X-Vault-Token': token} 16 | if cacerts: 17 | s.verify = cacerts 18 | 19 | click.echo('\nChecking Vault KV version...') 20 | kv_version = get_kv_version(client=s, host=host) 21 | click.echo(click.style(f'Using Vault KV version {kv_version}\n', bold=True, fg='white')) 22 | 23 | paths = [] 24 | with click.progressbar(range(count), label='Creating test keys in Vault') as bar: 25 | for _ in bar: 26 | path = common.key_path() 27 | if kv_version == 1: 28 | r = s.post(f'{host}/v1/secret/test/{path}', json={'value': common.random_data(size)}) 29 | else: 30 | r = s.post(f'{host}/v1/secret/data/test/{path}', json={'data': {'value': common.random_data(size)}}) 31 | 32 | if r.status_code >= 400: 33 | try: 34 | for msg in r.json()['warnings']: 35 | click.echo(click.style(f'Error returned by Vault: {msg}', bold=True, fg='yellow'), err=True) 36 | except KeyError: 37 | pass 38 | r.raise_for_status() 39 | else: 40 | paths.append(path) 41 | 42 | return paths 43 | 44 | 45 | @click.command() 46 | @click.option('--secrets', default=1000, 47 | help='Number of test secrets to create') 48 | @click.option('--secret-size', default=1024, 49 | help='Size of each secret, in bytes') 50 | @click.option('--transit-size', default=1048576, 51 | help='Size of data blocks to encrypt for Transit tests, in bytes') 52 | @click.option('--host', '-H', default=lambda: os.environ.get('VAULT_ADDR', 'http://localhost:8200'), 53 | help='URL of the Vault server to test') 54 | @click.option('--cacerts', default=lambda: os.environ.get('VAULT_CACERT', None), 55 | help='Path to a certificate chain for enabling TLS to Vault') 56 | @click.argument('token', envvar='VAULT_TOKEN') 57 | def main(host, secrets, secret_size, transit_size, cacerts, token): 58 | paths = populate(host, secrets, secret_size, cacerts, token) 59 | with open('testdata.json', 'w') as f: 60 | json.dump({ 61 | 'token': token, 62 | 'num_secrets': secrets, 63 | 'secret_size': secret_size, 64 | 'transit_size': transit_size, 65 | 'keys': paths, 66 | 'vault_cert': cacerts 67 | }, f) 68 | 69 | 70 | if __name__ == '__main__': 71 | main() 72 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | locustio==0.9.0 2 | click 3 | requests 4 | --------------------------------------------------------------------------------