v1.0.0

.gitignore

@@ -3,3 +3,4 @@
 /.direnv/
 
 *.json
+*.bin

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"
+}
+```

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=",

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;
   };
 }

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"];
+    };
+  };
+}

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"]

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='-'
+        )
     ))
 
 

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)