diff --git a/.gitignore b/.gitignore index d0176a5..231dc67 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,4 @@ /.direnv/ *.json +*.bin diff --git a/README.md b/README.md index e69de29..6135bd1 100644 --- a/README.md +++ b/README.md @@ -0,0 +1,18 @@ +# Simple Test Data Generator + +## Example Config + +```json +{ + "host": "127.0.0.1", + "port": 9250, + "log": "-", + "buffer-size": "4KiB", + "max-size": "2GB", + "keys": [ + "TESTKEY" + ], + "max-data": "10GB", + "database": "databse.json" +} +``` diff --git a/flake.lock b/flake.lock index 756f35b..9c06826 100644 --- a/flake.lock +++ b/flake.lock @@ -18,6 +18,46 @@ "type": "github" } }, + "flake-utils_2": { + "inputs": { + "systems": "systems_2" + }, + "locked": { + "lastModified": 1705309234, + "narHash": "sha256-uNRRNRKmJyCRC/8y1RqBkqWBLM034y4qN7EprSdmgyA=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "1ef2e671c3b0c19053962c07dbda38332dcebf26", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "microvm": { + "inputs": { + "flake-utils": "flake-utils", + "nixpkgs": [ + "nixpkgs" + ], + "spectrum": "spectrum" + }, + "locked": { + "lastModified": 1712654305, + "narHash": "sha256-CNdpLnGOUZfIhBanAFVF7t1xstaQGL4w6sQPrVeLlus=", + "owner": "astro", + "repo": "microvm.nix", + "rev": "ee0068ca87bdabbde3cc39b7af807c0302d0304c", + "type": "github" + }, + "original": { + "owner": "astro", + "repo": "microvm.nix", + "type": "github" + } + }, "nix-github-actions": { "inputs": { "nixpkgs": [ @@ -57,12 +97,12 @@ }, "poetry2nix-lib": { "inputs": { - "flake-utils": "flake-utils", + "flake-utils": "flake-utils_2", "nix-github-actions": "nix-github-actions", "nixpkgs": [ "nixpkgs" ], - "systems": "systems_2", + "systems": "systems_3", "treefmt-nix": "treefmt-nix" }, "locked": { @@ -81,10 +121,27 @@ }, "root": { "inputs": { + "microvm": "microvm", "nixpkgs": "nixpkgs", "poetry2nix-lib": "poetry2nix-lib" } }, + "spectrum": { + "flake": false, + "locked": { + "lastModified": 1708358594, + "narHash": "sha256-e71YOotu2FYA67HoC/voJDTFsiPpZNRwmiQb4f94OxQ=", + "ref": "refs/heads/main", + "rev": "6d0e73864d28794cdbd26ab7b37259ab0e1e044c", + "revCount": 614, + "type": "git", + "url": "https://spectrum-os.org/git/spectrum" + }, + "original": { + "type": "git", + "url": "https://spectrum-os.org/git/spectrum" + } + }, "systems": { "locked": { "lastModified": 1681028828, @@ -101,6 +158,21 @@ } }, "systems_2": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + }, + "systems_3": { "locked": { "lastModified": 1681028828, "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", diff --git a/flake.nix b/flake.nix index 558c8d4..54b7fd2 100644 --- a/flake.nix +++ b/flake.nix @@ -7,23 +7,32 @@ url = "github:nix-community/poetry2nix"; inputs.nixpkgs.follows = "nixpkgs"; }; + microvm = { + url = "github:astro/microvm.nix"; + inputs.nixpkgs.follows = "nixpkgs"; + }; }; - outputs = { - self, - nixpkgs, - poetry2nix-lib, - }: let + outputs = {self, ...} @ inputs: let supportedSystems = ["x86_64-linux" "x86_64-darwin" "aarch64-linux" "aarch64-darwin"]; - forAllSystems = nixpkgs.lib.genAttrs supportedSystems; - pkgs = forAllSystems (system: nixpkgs.legacyPackages.${system}); - poetry2nix = forAllSystems (system: poetry2nix-lib.lib.mkPoetry2Nix {pkgs = pkgs.${system};}); + forAllSystems = inputs.nixpkgs.lib.genAttrs supportedSystems; + pkgs = forAllSystems (system: inputs.nixpkgs.legacyPackages.${system}); + poetry2nix = forAllSystems (system: inputs.poetry2nix-lib.lib.mkPoetry2Nix {pkgs = pkgs.${system};}); in { # `nix build` packages = forAllSystems (system: { default = poetry2nix.${system}.mkPoetryApplication { projectDir = self; }; + vm = self.nixosConfigurations.vm.config.microvm.declaredRunner; + }); + + # `nix run` + apps = forAllSystems (system: { + default = { + program = "${self.packages.${system}.default}/bin/testdata"; + type = "app"; + }; }); # `nix fmt` @@ -45,5 +54,8 @@ ++ [poetryEnv]; }; }); + + # NixOS Module + nixosModules.default = import ./nix/module.nix inputs; }; } diff --git a/nix/module.nix b/nix/module.nix new file mode 100644 index 0000000..4fcb760 --- /dev/null +++ b/nix/module.nix @@ -0,0 +1,50 @@ +inputs: { + config, + lib, + pkgs, + system, + ... +}: let + cfg = config.testdata; + package = inputs.self.packages.${pkgs.stdenv.hostPlatform.system}.default; + inherit (lib) mkIf mkEnableOption mkOption types; + + format = pkgs.formats.json {}; + configFile = format.generate "config.json" cfg.settings; +in { + options.testdata = { + enable = mkEnableOption "testdata"; + + settings = mkOption { + type = with types; let + valueType = nullOr (oneOf [ + # TODO: restrict type to actual config file structure + bool + int + float + str + path + (attrsOf valueType) + (listOf valueType) + ]); + in + valueType; + default = throw "Please specify testdata.settings"; + }; + }; + + config = mkIf cfg.enable { + environment.systemPackages = [package]; + + systemd.services.testdata = { + enable = true; + + serviceConfig = { + Type = "simple"; + ExecStart = "${package}/bin/testdata --config ${configFile}"; + }; + + wantedBy = ["multi-user.target"]; + }; + }; +} diff --git a/pyproject.toml b/pyproject.toml index 58f23ea..94c4a1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ pydantic = "^2.6.4" ipaddress = "^1.0.23" [tool.poetry.scripts] -main = "src.main:main" +testdata = "src.main:main" [build-system] requires = ["poetry-core"] diff --git a/src/main.py b/src/main.py index 9a98e52..fedcdbd 100644 --- a/src/main.py +++ b/src/main.py @@ -2,90 +2,97 @@ import sys import asyncio import argparse import json +from os.path import exists -from fastapi import FastAPI, Request +from fastapi import FastAPI, Request, HTTPException from fastapi.responses import StreamingResponse from fastapi import status from hypercorn.config import Config from hypercorn.asyncio import serve import ipaddress +from .utils import convert_to_bytes, generate_data, load_database, save_database + # Setup Parser parser = argparse.ArgumentParser() parser.add_argument('-c', '--config', type=argparse.FileType('r'), - default='./config', help='Path to config file in JSON format.') -# parser.add_argument('-db', '--database', type=argparse.FileType('r'), # TODO: read+write -# default='./db.json', help='Path to database file in JSON format.') + default='./config.json', help='Path to config file in JSON format.') args = parser.parse_args(sys.argv[1:]) # Load Config -config = json.load(args.config) +CONFIG = json.load(args.config) +BUFFER_SIZE = convert_to_bytes(CONFIG['buffer-size']) +MAX_SIZE = convert_to_bytes(CONFIG['max-size']) +MAX_DATA = convert_to_bytes(CONFIG['max-data']) +AUTHORIZED_KEYS = CONFIG['keys'] +DATABASE = CONFIG['database'] + +if not exists(DATABASE): + save_database(DATABASE, {'data-used': 0}) api = FastAPI() -BUFFER_SIZE = 1024 * 4 # 4KB -HOST = '127.0.0.1' -PORT = 9250 + +class MaxSizePerRequestError(Exception): + pass -async def generate_test_data(size: int | str, max_size=1024 * 1024) -> bytes: - size_left = None +class MinSizePerRequestError(Exception): + pass - try: - size_left = int(size) - except ValueError: # treat as string - units = { - 'GB': 10 ** 9, 'GiB': 2 ** 30, - 'MB': 10 ** 6, 'MiB': 2 ** 20, - 'KB': 10 ** 3, 'KiB': 2 ** 10, - 'B': 1 - } - - for unit in units: - if size.endswith(unit): - size_left = int(float(size.removesuffix(unit)) * units[unit]) - break - - print('size_left', size_left) - - # yield data - while size_left > BUFFER_SIZE: - size_left -= BUFFER_SIZE - yield b'\0' * BUFFER_SIZE - else: - yield b'\0' * size_left - - -def check_policies(ip: str) -> None: - network = ipaddress.ip_network(ip) - print(network) @api.get('/') -async def test_data(size: str, request: Request) -> StreamingResponse: +def test_data(api_key: str, size: str, request: Request) -> StreamingResponse: try: - check_policies(request.client.host) + if api_key not in AUTHORIZED_KEYS: + raise HTTPException( + status_code=status.HTTP_401_UNAUTHORIZED, + detail='Invalid API Key.' + ) + + size = convert_to_bytes(size) + + if size < 0: + raise MinSizePerRequestError + elif MAX_SIZE < size: + raise MaxSizePerRequestError + + database = load_database(DATABASE) + if MAX_DATA <= database['data-used'] + size: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail='Service not available.' + ) + database['data-used'] += size + + save_database(DATABASE, database) return StreamingResponse( status_code=status.HTTP_200_OK, - content=generate_test_data(size) + content=generate_data(size, BUFFER_SIZE) + ) + + except MinSizePerRequestError: + raise HTTPException( + status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE, + detail=f'Size has to be not-negative.' + ) + except MaxSizePerRequestError: + raise HTTPException( + status_code=status.HTTP_416_REQUESTED_RANGE_NOT_SATISFIABLE, + detail=f'Exceeded max size per request of {MAX_SIZE} Bytes.' ) - # except TooManyRequestsError as err: - # ... - # except BlockedError as err: - # ... - # except OutOfQuotaError as err: - # ... - except err: - raise err - pass def main(): asyncio.run(serve( api, - Config().from_object({'bind': [f'{HOST}:{PORT}']}) + Config().from_mapping( + bind=[f"{CONFIG['host']}:{CONFIG['port']}"], + accesslog='-' + ) )) diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000..f7ad115 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,43 @@ +import json + +def convert_to_bytes(size: int | str) -> int: + try: + return int(size) + except ValueError: # treat as string + units = { + 'TB': 10 ** 12, 'TiB': 2 ** 40, + 'GB': 10 ** 9, 'GiB': 2 ** 30, + 'MB': 10 ** 6, 'MiB': 2 ** 20, + 'KB': 10 ** 3, 'KiB': 2 ** 10, + 'B': 1 + } + + for unit in units: + if size.endswith(unit): + return int(float(size.removesuffix(unit)) * units[unit]) + break + + +async def generate_data(size: int, buffer_size: int = 4 * 1024) -> bytes: + size_left = size + + while size_left > buffer_size: + size_left -= buffer_size + yield b'\0' * buffer_size + else: + yield b'\0' * size_left + + +def check_policies(ip: str) -> None: + network = ipaddress.ip_network(ip) + print(network) + + +def load_database(path: str) -> dict: + with open(path, 'r') as file: + return json.load(file) + + +def save_database(path: str, database: dict) -> None: + with open(path, 'w') as file: + json.dump(database, file, indent=2)