admin/trailofbits-dropkit
1
0
Fork 0
mirror of https://github.com/trailofbits/dropkit synced 2026-06-21 14:11:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
Dan Guido
2026-01-29 21:50:14 -05:00
commit bc36073036
38 changed files with 12434 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.claude/settings.local.json
+19
View File
@@ -0,0 +1,19 @@
{
"permissions": {
"allow": [
"Bash(uv sync:*)",
"Bash(uv run pytest:*)",
"Bash(uv add:*)",
"Bash(uv run ruff:*)",
"Bash(uv run mypy:*)",
"Bash(./lint.sh:*)",
"Bash(uv run tobcloud --help:*)",
"Bash(uv run tobcloud on --help:*)",
"Bash(uv run tobcloud off --help:*)",
"Bash(git add:*)",
"Bash(git commit:*)"
],
"deny": [],
"ask": []
}
}
.claude/skills/git-worktrees/SKILL.md
+83
View File
@@ -0,0 +1,83 @@
---
name: git-worktrees
description: Use git worktrees when running multiple Claude Code instances in parallel for different features - creates isolated workspaces with separate branches and virtual environments
---
# Parallel Development with Git Worktrees
## Overview
Git worktrees create isolated workspaces sharing the same repository, allowing work on multiple branches simultaneously without switching. This is essential when running multiple Claude Code instances in parallel.
**Core principle:** One worktree per Claude Code instance ensures isolation and prevents conflicts.
**Announce at start:** "I'm using the git-worktrees skill to set up an isolated workspace."
## When to Use
- Running multiple Claude Code instances for different features
- Need to work on a feature branch while keeping main branch accessible
- Parallel development on separate tasks
## Setup
```bash
# Create worktrees for parallel development
git worktree add ../dropkit-feature-a feature-a
git worktree add ../dropkit-feature-b feature-b
# Each worktree gets its own directory with full codebase
# Run Claude Code in each directory independently
```
## Guidelines for Parallel Instances
1. **One worktree per Claude Code instance** - Never run multiple instances in the same directory
2. **Separate branches** - Each worktree should be on its own feature branch
3. **Independent `uv sync`** - Run `uv sync` in each worktree (creates separate `.venv`)
4. **Tests run independently** - Each worktree can run its own test suite without conflicts
5. **Merge via main branch** - When features are complete, merge branches to main
## Managing Worktrees
```bash
# List all worktrees
git worktree list
# Remove a worktree when done
git worktree remove ../dropkit-feature-a
# Prune stale worktree entries
git worktree prune
```
## Potential Conflicts to Avoid
| Resource | Risk | Mitigation |
|----------|------|------------|
| `~/.config/dropkit/` | User config is shared | Don't modify during parallel dev |
| `~/.ssh/config` | SSH config is shared | Coordinate droplet names |
| DigitalOcean API | Creating droplets with same name | Use unique droplet names per worktree |
## Quick Reference
| Situation | Action |
|-----------|--------|
| Starting parallel work | Create new worktree with feature branch |
| New worktree created | Run `uv sync` to create isolated venv |
| Feature complete | Merge to main, remove worktree |
| Stale worktrees | Run `git worktree prune` |
## Example Workflow
```
You: I'm using the git-worktrees skill to set up an isolated workspace.
[Create worktree: git worktree add ../dropkit-auth feature/auth]
[Run uv sync]
[Run uv run pytest - all passing]
Worktree ready at ../dropkit-auth
Tests passing
Ready to implement auth feature
```
.coverage
BIN
View File
Binary file not shown.
.devcontainer/Dockerfile
+29
View File
@@ -0,0 +1,29 @@
FROM mcr.microsoft.com/devcontainers/base:debian
SHELL ["/bin/bash", "-o", "pipefail", "-c"]
# Install uv
COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /usr/local/bin/
# Switch to non-root user
USER vscode
WORKDIR /home/vscode
# Add ~/.local/bin to PATH for Claude Code
ENV PATH="/home/vscode/.local/bin:$PATH"
# Install Claude Code (installs to ~/.local/bin/claude)
RUN curl -fsSL https://claude.ai/install.sh | bash && \
claude plugin marketplace add trailofbits/skills && \
claude plugin marketplace add anthropics/skills && \
claude plugin marketplace add obra/superpowers
# Set environment variables
ENV DEVCONTAINER=true
ENV EDITOR=nano
# Configure zsh and add aliases
RUN echo "export HISTFILE=\"\$HOME/.zsh_history_data/.zsh_history\"" >> ~/.zshrc && \
echo "alias yolo=\"claude --dangerously-skip-permissions\"" >> ~/.zshrc
WORKDIR /workspace
.devcontainer/devcontainer.json
+58
View File
@@ -0,0 +1,58 @@
{
"$schema": "https://raw.githubusercontent.com/devcontainers/spec/main/schemas/devContainer.schema.json",
"name": "dropkit",
"build": {
"dockerfile": "Dockerfile"
},
"runArgs": ["--hostname=dropkit-dev"],
"features": {
"ghcr.io/devcontainers/features/python:1": {
"version": "3.13"
},
"ghcr.io/devcontainers/features/github-cli:1": {},
"ghcr.io/tailscale/codespace/tailscale:1": {}
},
"remoteUser": "vscode",
"workspaceFolder": "/workspace",
"workspaceMount": "source=${localWorkspaceFolder},target=/workspace,type=bind,consistency=cached",
"mounts": [
"source=dropkit-zsh-history,target=/home/vscode/.zsh_history_data,type=volume",
"source=dropkit-claude,target=/home/vscode/.claude_data,type=volume",
"source=dropkit-gh,target=/home/vscode/.gh_data,type=volume",
"source=dropkit-config,target=/home/vscode/.config/dropkit,type=volume",
"source=dropkit-ssh,target=/home/vscode/.ssh,type=volume"
],
"customizations": {
"vscode": {
"extensions": [
"anthropic.claude-code",
"ms-python.python",
"ms-python.vscode-pylance",
"charliermarsh.ruff",
"eamodio.gitlens",
"tamasfe.even-better-toml"
],
"settings": {
"terminal.integrated.defaultProfile.linux": "zsh",
"python.defaultInterpreterPath": "/workspace/.venv/bin/python",
"python.terminal.activateEnvironment": true,
"[python]": {
"editor.formatOnSave": true,
"editor.defaultFormatter": "charliermarsh.ruff",
"editor.codeActionsOnSave": {
"source.fixAll.ruff": "explicit",
"source.organizeImports.ruff": "explicit"
}
}
}
}
},
"onCreateCommand": "uv sync",
"updateContentCommand": "uv sync",
"postCreateCommand": "sudo chown -R vscode:vscode /home/vscode/.config/dropkit /home/vscode/.claude_data /home/vscode/.gh_data /home/vscode/.zsh_history_data /home/vscode/.ssh",
"containerEnv": {
"UV_LINK_MODE": "copy",
"CLAUDE_CONFIG_DIR": "/home/vscode/.claude_data",
"GH_CONFIG_DIR": "/home/vscode/.gh_data"
}
}
.github/ISSUE_TEMPLATE/bug_report.md
+38
View File
@@ -0,0 +1,38 @@
---
name: Bug report
about: Create a report to help us improve
title: ''
labels: ''
assignees: ''
---
**Describe the bug**
A clear and concise description of what the bug is.
**To Reproduce**
Steps to reproduce the behavior:
1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error
**Expected behavior**
A clear and concise description of what you expected to happen.
**Screenshots**
If applicable, add screenshots to help explain your problem.
**Desktop (please complete the following information):**
- OS: [e.g. iOS]
- Browser [e.g. chrome, safari]
- Version [e.g. 22]
**Smartphone (please complete the following information):**
- Device: [e.g. iPhone6]
- OS: [e.g. iOS8.1]
- Browser [e.g. stock browser, safari]
- Version [e.g. 22]
**Additional context**
Add any other context about the problem here.
.github/ISSUE_TEMPLATE/feature_request.md
+20
View File
@@ -0,0 +1,20 @@
---
name: Feature request
about: Suggest an idea for this project
title: ''
labels: ''
assignees: ''
---
**Is your feature request related to a problem? Please describe.**
A clear and concise description of what the problem is. Ex. I'm always frustrated when [...]
**Describe the solution you'd like**
A clear and concise description of what you want to happen.
**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.
**Additional context**
Add any other context or screenshots about the feature request here.
.github/dependabot.yml
+23
View File
@@ -0,0 +1,23 @@
version: 2
updates:
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
cooldown:
default-days: 7 # Supply chain protection: wait for community vetting
groups:
github-actions:
patterns:
- "*"
- package-ecosystem: "pip"
directory: "/"
schedule:
interval: "weekly"
cooldown:
default-days: 7 # Supply chain protection: wait for community vetting
groups:
python-dependencies:
patterns:
- "*"
.github/workflows/ci.yml
+45
View File
@@ -0,0 +1,45 @@
name: CI
on:
push:
branches: [main]
pull_request:
branches: [main]
permissions:
contents: read
jobs:
test:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.11", "3.12", "3.13"]
steps:
- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- name: Install uv
uses: astral-sh/setup-uv@0c5e2b8115b80b4c7c5ddf6ffdd634974642d182 # v5.4.1
with:
enable-cache: true
- name: Set up Python ${{ matrix.python-version }}
run: uv python install ${{ matrix.python-version }}
- name: Install dependencies
run: uv sync
- name: Run ruff check
run: uv run ruff check .
- name: Run ruff format check
run: uv run ruff format --check .
- name: Run type check
run: uv run ty check dropkit/
- name: Run tests
run: uv run pytest
.gitignore
+13
View File
@@ -0,0 +1,13 @@
# Python-generated files
__pycache__/
*.py[oc]
build/
dist/
wheels/
*.egg-info
# Build-generated version file
dropkit/_version.txt
# Virtual environments
.venv
.pre-commit-config.yaml
+42
View File
@@ -0,0 +1,42 @@
# Pre-commit configuration
# Run with: prek run (or prek run --all-files)
default_language_version:
python: python3.11
repos:
# Ruff - official pre-commit hooks (faster than local)
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.14.14
hooks:
- id: ruff-check
args: [--fix]
- id: ruff-format
# prek builtin hooks - fast, no external deps
- repo: builtin
hooks:
- id: trailing-whitespace
- id: end-of-file-fixer
- id: check-yaml
exclude: ^dropkit/templates/ # Jinja2 templates aren't valid YAML
- id: check-toml
- id: check-merge-conflict
- id: check-added-large-files
# ty type checking (no official pre-commit hook, use local)
- repo: local
hooks:
- id: ty-check
name: ty type check
entry: uv run ty check dropkit/
language: system
types: [python]
pass_filenames: false
# Shell script linting
- repo: https://github.com/koalaman/shellcheck-precommit
rev: v0.10.0
hooks:
- id: shellcheck
args: [--severity=error]
.python-version
+1
View File
@@ -0,0 +1 @@
3.11
CLAUDE.md
+171
View File
@@ -0,0 +1,171 @@
# dropkit
CLI tool for managing DigitalOcean droplets for Trail of Bits engineers.
Pre-configured cloud-init, Tailscale VPN (enabled by default), and SSH config management.
## Critical Rules ⚠️
- **Never use `pip`** — always use `uv` for all Python operations
- **Always run `prek run`** before committing (or `prek install` to auto-run on commit)
- **Keep README.md in sync** when adding commands or features
## Quick Commands
```bash
uv sync # Install dependencies
prek install # Set up pre-commit hooks (one-time)
prek run # Run all checks (ruff, ty, shellcheck, etc.)
uv run pytest # Run tests
uv run dropkit --help # CLI help
```
## Project Structure
```
dropkit/
├── pyproject.toml # Dependencies and metadata
├── CLAUDE.md # This file
├── README.md # User documentation
├── dropkit/
│ ├── main.py # Typer CLI entry point
│ ├── config.py # Config with SSH key validation
│ ├── api.py # DigitalOcean REST API (see docstring for endpoints)
│ ├── cloudinit.py # Cloud-init template rendering
│ ├── ssh_config.py # SSH config manipulation
│ └── templates/
│ └── default-cloud-init.yaml
└── tests/
├── test_api.py
├── test_config.py
├── test_main_helpers.py
└── test_ssh_config.py
```
## Technology Stack
- **Python 3.11+** with `uv` (NOT pip)
- **CLI**: Typer + Rich
- **API**: Direct REST calls (requests library, no SDK)
- **Config**: YAML + Pydantic 2.x validation
- **Templating**: Jinja2 for cloud-init
- **Code Quality**: Ruff (linter + formatter), ty (types)
## Key Conventions
### Username
- **Derived from DigitalOcean account email**, not configured
- Fetched via `/v2/account`, sanitized for Linux compatibility
- `john.doe@trailofbits.com` → `john_doe`
### SSH Hostname
- All SSH entries use `dropkit.<droplet-name>` format
- Centralized in `get_ssh_hostname()` helper
### Tags
- Default tags: `owner:<username>` and `firewall`
- Additional tags **extend** defaults (never replace)
- Used for filtering, billing, and security
### SSH Keys
- Only public keys (*.pub) accepted
- Strict validation rejects private keys
- Auto-detects id_ed25519.pub, id_rsa.pub, id_ecdsa.pub
### Tailscale VPN
- **Enabled by default** for new droplets
- Locks down UFW to only allow tailscale0 interface
- Disable with `--no-tailscale` flag
## Architecture Decisions
1. **Username from email** — Not stored in config; fetched from DO API on demand.
Ensures consistency across machines.
2. **SSH key validation** — Prevents accidental private key upload.
Validates format (ssh-rsa, ssh-ed25519, ecdsa-sha2-*, etc.).
3. **Tags extend defaults** — `--tags` adds to defaults, never replaces.
Ensures owner tag always present.
4. **Direct REST API** — No python-digitalocean library.
Simpler, fewer dependencies, full control.
5. **Tailscale by default** — Secure VPN access without public SSH.
Local Tailscale required; keeps public IP if local Tailscale not running.
6. **Pydantic validation** — Runtime type safety for config files.
Clear errors for invalid configurations.
7. **SSH config backups** — Created at `~/.ssh/config.bak` before modifications.
Each backup overwrites previous.
## Gotchas & Troubleshooting
### Cloud-init JSON parsing (CRITICAL)
`cloud-init status --format=json` may return **non-zero exit code but valid JSON**.
**Always parse JSON regardless of subprocess return code.**
```python
# CORRECT: Parse JSON even on error
result = subprocess.run([...], capture_output=True)
data = json.loads(result.stdout) # Don't check returncode first
# WRONG: Checking returncode before parsing
if result.returncode == 0: # May skip valid JSON with error status
data = json.loads(result.stdout)
```
Status values: `"done"` (success), `"error"` (failed), `"running"` (in progress).
### Disk resize is permanent
Cannot be undone. Use `--no-disk` to resize only CPU/memory.
### SSH config backups overwrite
Only keeps one backup at `~/.ssh/config.bak`.
## Development
### Package Management
```bash
uv add <package> # Add dependency
uv add --dev <package> # Add dev dependency
uv sync # Install all
uv run <command> # Run in venv
```
### Linting
```bash
prek run # Run all checks (required before commit)
prek run --all-files # Check all files, not just staged
uv run ruff check --fix . # Lint + autofix only
uv run ty check dropkit/ # Type check only
```
**Ruff config**: Python 3.11+, 100-char lines, modern syntax (`str | None`, `list[str]`).
### Testing
```bash
uv run pytest # All tests
uv run pytest tests/test_api.py # Specific file
uv run pytest -k "validate_ssh" # Pattern match
uv run pytest -v # Verbose
```
## Pydantic Models
- **`DropkitConfig`** — Root config with `extra='forbid'`
- **`DigitalOceanConfig`** — API token validation
- **`DefaultsConfig`** — Region, size, image slugs
- **`CloudInitConfig`** — Template path, SSH keys (min 1)
- **`SSHConfig`** — SSH config path, identity file
- **`TailscaleConfig`** — VPN settings (enabled, lock_down_firewall, auth_timeout)
Config files: `~/.config/dropkit/config.yaml`, `~/.config/dropkit/cloud-init.yaml`
## Shell Completion
```bash
dropkit --install-completion zsh # Enable tab completion
```
Provides dynamic completion for droplet names (filtered by owner tag).
LICENSE
+190
View File
@@ -0,0 +1,190 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to the Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
Copyright 2025 Trail of Bits, Inc.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
Makefile
+16
View File
@@ -0,0 +1,16 @@
.PHONY: dev lint format test audit
dev:
uv sync --all-groups
lint:
uv run ruff format --check . && uv run ruff check . && uv run ty check dropkit/
format:
uv run ruff format .
test:
uv run pytest
audit:
uv run pip-audit
README.md
+259
View File
@@ -0,0 +1,259 @@
# dropkit
A command-line tool for managing DigitalOcean droplets with automated setup, SSH configuration, and lifecycle management.
## Features
- 🚀 **Quick droplet creation** with cloud-init automation
- 🔑 **Automatic SSH configuration** - just run `ssh dropkit.<droplet-name>`
- 🔐 **Tailscale VPN** - secure access via Tailscale (enabled by default)
- 👤 **User management** - automatically creates your user account on droplets
- 🏷️ **Smart tagging** - organizes droplets by owner for easy filtering
- 🔒 **Security-first** - SSH key validation, confirmation prompts for destructive operations
- 📊 **Rich CLI** - beautiful tables, progress indicators, and helpful error messages
- 🔄 **Complete lifecycle** - create, list, resize, destroy droplets with ease
- 💤 **Hibernate/Wake** - snapshot and destroy to save costs, restore later with one command
## Prerequisites
- **Python 3.11+**
- **DigitalOcean account** with an API token ([create one here](https://cloud.digitalocean.com/account/api/tokens))
- **SSH key pair** (usually `~/.ssh/id_ed25519.pub` or `~/.ssh/id_rsa.pub`)
- **uv** package manager ([install instructions](https://github.com/astral-sh/uv))
- **Tailscale** (optional but recommended) - install from [tailscale.com/download](https://tailscale.com/download)
## Installation
```bash
# HTTPS or SSH depending on your GitHub setup
uv tool install git+https://github.com/trailofbits/dropkit.git
uv tool install git+ssh://git@github.com/trailofbits/dropkit.git
# Upgrade
uv tool upgrade dropkit
```
## Quick Start
### 1. Initialize Configuration
Run the initialization wizard:
```bash
dropkit init
```
This will validate your DigitalOcean API token, detect SSH keys, register them with DigitalOcean, and let you choose defaults (type `?` for help to see available options).
### 2. Create Your First Droplet
```bash
# Interactive mode - prompts for name, region, size, image (type ? for help)
dropkit create
# Or specify the name and use defaults
dropkit create my-first-droplet
# Assign to a specific project (by name or ID)
dropkit create my-droplet --project "My Project"
# Create without Tailscale VPN
dropkit create my-droplet --no-tailscale
```
The tool will:
1. Create the droplet and wait for it to become active
2. Add SSH configuration automatically
3. Wait for cloud-init to complete
4. **Tailscale setup** (enabled by default):
- Display an auth URL for you to authenticate in your browser
- Update SSH config with your Tailscale IP
- Lock down the firewall to only allow Tailscale traffic
### 3. Connect via SSH
```bash
ssh dropkit.my-first-droplet
```
Your user account is already set up with your SSH keys.
## Available Commands
```
Usage: dropkit [OPTIONS] COMMAND [ARGS]...
Manage DigitalOcean droplets
Commands:
init Initialize dropkit configuration.
create Create a new DigitalOcean droplet with cloud-init configuration.
list List droplets and hibernated snapshots tagged with owner:<username>.
config-ssh Configure SSH for an existing droplet.
info Show detailed information about a droplet.
rename Rename a droplet (requires confirmation).
destroy Destroy a droplet or hibernated snapshot (DESTRUCTIVE).
resize Resize a droplet (causes downtime - requires power off).
on Power on a droplet.
off Power off a droplet (requires confirmation).
hibernate Hibernate a droplet (snapshot and destroy to save costs).
wake Wake a hibernated droplet (restore from snapshot).
enable-tailscale Enable Tailscale VPN on an existing droplet.
list-ssh-keys List SSH keys registered via dropkit.
add-ssh-key Add or import an SSH public key to DigitalOcean.
delete-ssh-key Delete an SSH key registered via dropkit.
version Show the version of dropkit.
```
Use `dropkit <command> --help` for detailed help on any command.
## Configuration
Configuration files are stored in `~/.config/dropkit/`:
- **`config.yaml`** - Main configuration (API token, defaults, SSH keys)
- **`cloud-init.yaml`** - Cloud-init template (customizable)
### Default Tags
All droplets are automatically tagged with:
- `owner:<username>` - Your DigitalOcean account username (derived from email)
- `firewall` - For security group identification
### Projects
- **Set default project** during `dropkit init` (type `?` to see available projects)
- **Override per-droplet** using `--project <name>` with `dropkit create`
- Specify by name or UUID; tab completion available
### SSH Hostname Convention
All SSH config entries use the prefix `dropkit.<droplet-name>`:
- Connect with: `ssh dropkit.my-droplet`
### Shell Completion
Enable tab completion for droplet names in your shell:
**Zsh (recommended):**
```bash
dropkit --install-completion zsh
```
**Bash:**
```bash
dropkit --install-completion bash
```
After installation, restart your shell. Tab completion dynamically fetches your droplets from DigitalOcean:
```bash
dropkit info <TAB> # Shows your droplets
dropkit destroy <TAB> # Shows your droplets
dropkit resize <TAB> # Shows your droplets
dropkit on <TAB> # Shows your droplets
dropkit off <TAB> # Shows your droplets
```
### Hibernate and Wake (Cost Saving)
DigitalOcean charges for stopped droplets at the full hourly rate. To avoid this, use hibernate/wake:
```bash
# Hibernate: snapshot the droplet and destroy it (stops billing)
dropkit hibernate my-droplet
# Wake: restore the droplet from the snapshot
dropkit wake my-droplet
# Delete a hibernated snapshot without restoring
dropkit destroy my-droplet
```
**How it works:**
1. `hibernate` powers off the droplet, creates a snapshot (`dropkit-<name>`), then destroys the droplet
2. `wake` creates a new droplet from the snapshot with the same region and size
3. Snapshots are tagged with `owner:<username>` and `size:<size-slug>` for tracking
4. After waking, you're prompted to delete the snapshot (default: yes)
**Note:** Snapshots are billed at $0.06/GB/month, which is typically much cheaper than keeping a droplet running.
### Cloud-Init Customization
Edit `~/.config/dropkit/cloud-init.yaml` to customize user setup, package installation, firewall rules, and shell configuration. The template uses Jinja2 syntax with variables `{{ username }}` and `{{ ssh_keys }}`.
## Troubleshooting
### "Config not found. Run 'dropkit init' first"
Initialize the configuration:
```bash
dropkit init
```
### Cloud-init failed or timeout
Check cloud-init status manually:
```bash
ssh dropkit.my-droplet 'sudo cloud-init status'
ssh dropkit.my-droplet 'sudo cat /var/log/cloud-init.log'
```
Use `--verbose` flag to see detailed output:
```bash
dropkit create my-droplet --verbose
```
### "Droplet not found with tag owner:<username>"
The droplet might belong to someone else. List your droplets:
```bash
dropkit list
```
## Technology Stack
- **CLI Framework**: [Typer](https://typer.tiangolo.com/) - Modern CLI framework
- **UI/Display**: [Rich](https://rich.readthedocs.io/) - Terminal formatting
- **API Client**: [requests](https://requests.readthedocs.io/) - HTTP library
- **Configuration**: [Pydantic](https://docs.pydantic.dev/) - Data validation
- **Templating**: [Jinja2](https://jinja.palletsprojects.com/) - Cloud-init templates
- **Package Manager**: [uv](https://github.com/astral-sh/uv) - Fast Python package manager
- **Code Quality**: Ruff (linter/formatter) + ty (type checker)
## Appendix: API Token Permissions
### Creating Your Token
1. Go to [DigitalOcean API Tokens](https://cloud.digitalocean.com/account/api/tokens)
2. Click **Generate New Token** with name "dropkit-cli"
3. Select **Custom Scopes** (recommended) or **Full Access** (simpler)
4. For custom scopes, enable the 23 scopes listed below
### Required Scopes (23 total)
`account:read`, `actions:read`, `droplet:create`, `droplet:read`, `droplet:update`, `droplet:delete`, `image:create`, `image:read`, `image:update`, `image:delete`, `project:read`, `project:update`, `regions:read`, `sizes:read`, `snapshot:read`, `snapshot:delete`, `ssh_key:create`, `ssh_key:read`, `ssh_key:update`, `ssh_key:delete`, `tag:read`, `tag:create`, `vpc:read`
### Scope Reference by Feature
| Feature | Required Scopes |
|---------|----------------|
| **Initialize config** | `account:read`, `regions:read`, `sizes:read`, `image:read`, `ssh_key:read`, `ssh_key:create`, `project:read` |
| **Create droplets** | `droplet:create`, `project:update`, `actions:read`, `tag:create` |
| **List droplets** | `droplet:read`, `snapshot:read`, `tag:read` |
| **Show droplet info** | `droplet:read` |
| **Destroy droplets** | `droplet:delete`, `snapshot:delete` |
| **Rename droplets** | `droplet:update` |
| **Resize droplets** | `droplet:update`, `sizes:read`, `actions:read` |
| **Power on/off** | `droplet:update`, `actions:read` |
| **Hibernate** | `droplet:update`, `droplet:delete`, `snapshot:create`, `actions:read`, `tag:create` |
| **Wake** | `droplet:create`, `snapshot:read`, `snapshot:delete` |
| **Manage SSH keys** | `ssh_key:read`, `ssh_key:create`, `ssh_key:update`, `ssh_key:delete` |
For more information, see the [DigitalOcean API Token Scopes documentation](https://docs.digitalocean.com/reference/api/scopes/).
dropkit/__init__.py
+54
View File
@@ -0,0 +1,54 @@
"""dropkit - Manage DigitalOcean droplets for ToB engineers."""
import subprocess
from pathlib import Path
def _get_version() -> str:
"""
Get version string.
Returns:
- "dev" if running from git repository (development mode)
- "0.1.0+git.<commit>" if installed (commit hash embedded at build time)
- "0.1.0" fallback if version cannot be determined
"""
# Check if we're in a git repository (development mode)
try:
repo_root = Path(__file__).parent.parent
if (repo_root / ".git").exists():
return "dev"
except Exception:
pass
# Installed mode - check for embedded version file (created at build time)
try:
version_file = Path(__file__).parent / "_version.txt"
if version_file.exists():
commit = version_file.read_text().strip()
if commit:
return f"0.1.0+git.{commit}"
except Exception:
pass
# Fallback: try to get commit from git (in case running from source)
try:
result = subprocess.run(
["git", "rev-parse", "--short=7", "HEAD"],
capture_output=True,
text=True,
timeout=1,
cwd=Path(__file__).parent,
)
if result.returncode == 0:
commit = result.stdout.strip()
if commit:
return f"0.1.0+git.{commit}"
except Exception:
pass
# Final fallback to base version
return "0.1.0"
__version__ = _get_version()
dropkit/api.py
+954
View File
@@ -0,0 +1,954 @@
"""DigitalOcean API client using raw REST API calls.
API Reference
=============
Base URL: https://api.digitalocean.com/v2
Auth: Authorization: Bearer <token>
Key Endpoints
-------------
Account & SSH Keys:
GET /account Account info (includes email for username)
GET /account/keys List SSH keys (paginated)
GET /account/keys/{fingerprint} Get SSH key by fingerprint
POST /account/keys Add new SSH key
PUT /account/keys/{id} Update SSH key name
DELETE /account/keys/{id} Delete SSH key
Droplets:
POST /droplets Create droplet
GET /droplets List droplets
GET /droplets?tag_name=X Filter by tag
GET /droplets/{id} Get droplet info
DELETE /droplets/{id} Delete droplet
POST /droplets/{id}/actions Perform action (resize, power_on, power_off, snapshot)
Metadata:
GET /regions List regions (paginated)
GET /sizes List droplet sizes (paginated)
GET /images List images (paginated)
Actions:
GET /actions/{id} Check action status
Projects:
GET /projects List projects (paginated)
GET /projects/{project_id} Get project by UUID
GET /projects/default Get default project
POST /projects/{project_id}/resources Assign resources (body: {"resources": ["do:droplet:123"]})
Snapshots:
GET /snapshots List snapshots (filter: resource_type=droplet)
GET /snapshots/{id} Get snapshot
DELETE /snapshots/{id} Delete snapshot
Tags:
POST /tags Create tag
POST /tags/{name}/resources Tag a resource
Pagination
----------
Uses `page` and `per_page` query params (max 200/page).
This module auto-handles pagination by following `links.pages.next` URLs.
"""
import re
from typing import Any
import requests
class DigitalOceanAPIError(Exception):
"""Exception raised for DigitalOcean API errors."""
def __init__(self, message: str, status_code: int | None = None):
self.status_code = status_code
super().__init__(message)
class DigitalOceanAPI:
"""Client for DigitalOcean REST API."""
def __init__(self, token: str):
"""Initialize API client with authentication token."""
self.token = token
self.base_url = "https://api.digitalocean.com/v2"
self.session = requests.Session()
self.session.headers.update(
{
"Authorization": f"Bearer {token}",
"Content-Type": "application/json",
}
)
def get_account(self) -> dict[str, Any]:
"""
Get account information.
Returns:
Account object with details like email, droplet_limit, status, etc.
"""
response = self._request("GET", "/account")
return response.get("account", {})
def get_username(self) -> str:
"""
Get username from DigitalOcean account email.
Fetches the account information and sanitizes the email address
to create a valid Linux username.
Returns:
Sanitized username from DigitalOcean account
Raises:
DigitalOceanAPIError: If account email cannot be fetched or is empty
"""
account = self.get_account()
email = account.get("email", "")
if not email:
raise DigitalOceanAPIError("No email found in DigitalOcean account")
return self._sanitize_email_for_username(email)
@staticmethod
def _sanitize_email_for_username(email: str) -> str:
"""
Sanitize email address to create a valid username.
Removes @trailofbits.com suffix and replaces special characters.
Args:
email: Email address from DigitalOcean account
Returns:
Sanitized username suitable for Linux user creation
"""
# Remove @trailofbits.com suffix (case insensitive)
username = re.sub(r"@trailofbits\.com$", "", email, flags=re.IGNORECASE)
# If no @trailofbits.com, just take the part before @
if "@" in username:
username = username.split("@")[0]
# Replace dots, hyphens, and other special characters with underscores
username = re.sub(r"[^a-z0-9_]", "_", username.lower())
# Remove leading/trailing underscores
username = username.strip("_")
# Ensure it starts with a letter (Linux username requirement)
if username and not username[0].isalpha():
username = "u" + username
# Fallback if sanitization results in empty string
if not username:
username = "user"
return username
@staticmethod
def _validate_positive_int(value: int, name: str) -> None:
"""
Validate that an integer ID is positive.
Args:
value: The integer to validate
name: Name of the parameter (for error message)
Raises:
ValueError: If the value is not positive
"""
if value <= 0:
raise ValueError(f"{name} must be a positive integer, got: {value}")
def _request(
self,
method: str,
endpoint: str,
**kwargs,
) -> dict[str, Any]:
"""Make an API request."""
url = f"{self.base_url}{endpoint}"
try:
response = self.session.request(method, url, **kwargs)
response.raise_for_status()
if response.status_code == 204 and response.text == "":
return {}
return response.json()
except requests.exceptions.HTTPError as e:
error_msg = f"API request failed: {e}"
if e.response is not None:
try:
error_data = e.response.json()
if "message" in error_data:
error_msg = f"API error: {error_data['message']}"
except ValueError:
pass
raise DigitalOceanAPIError(error_msg, e.response.status_code)
raise DigitalOceanAPIError(error_msg) from e
except requests.exceptions.RequestException as e:
raise DigitalOceanAPIError(f"Network error: {e}") from e
def _get_paginated(
self,
endpoint: str,
key: str,
per_page: int = 200,
extra_params: dict[str, str] | None = None,
max_pages: int = 1000,
) -> list[dict[str, Any]]:
"""
Fetch all pages of a paginated API endpoint.
Args:
endpoint: API endpoint to fetch
key: Key in response containing the items (e.g., 'regions', 'sizes')
per_page: Number of items per page (default 200, max 200)
extra_params: Additional query parameters to include (safely encoded)
max_pages: Maximum number of pages to fetch (default 1000, prevents DoS)
Returns:
List of all items across all pages
Raises:
DigitalOceanAPIError: If max_pages limit is reached
"""
all_items = []
page = 1
while True:
# Build params dict safely
params: dict[str, str | int] = {"page": page, "per_page": per_page}
if extra_params:
params.update(extra_params)
response = self._request("GET", endpoint, params=params)
items = response.get(key, [])
all_items.extend(items)
# Check if there are more pages
links = response.get("links", {})
pages = links.get("pages", {})
# If there's no "next" link, we're done
if "next" not in pages:
break
# Safety: prevent infinite pagination
if page >= max_pages:
raise DigitalOceanAPIError(
f"Pagination limit reached: {max_pages} pages. "
"This may indicate an API issue or misconfiguration."
)
page += 1
return all_items
def get_regions(self) -> list[dict[str, Any]]:
"""
Fetch all available regions (handles pagination).
Returns:
List of region objects with slug, name, available, features, etc.
"""
return self._get_paginated("/regions", "regions")
def get_sizes(self) -> list[dict[str, Any]]:
"""
Fetch all available droplet sizes (handles pagination).
Returns:
List of size objects with slug, memory, vcpus, disk, price, etc.
"""
return self._get_paginated("/sizes", "sizes")
def get_available_regions(self) -> list[dict[str, Any]]:
"""Get only available regions."""
regions = self.get_regions()
return [r for r in regions if r.get("available", False)]
def get_available_sizes(self) -> list[dict[str, Any]]:
"""Get only available sizes."""
sizes = self.get_sizes()
return [s for s in sizes if s.get("available", False)]
def get_images(self, image_type: str = "distribution") -> list[dict[str, Any]]:
"""
Fetch all available images (handles pagination).
Args:
image_type: Filter by image type ('distribution', 'application', or 'all')
Returns:
List of image objects with slug, name, distribution, etc.
"""
if image_type == "all":
return self._get_paginated("/images", "images")
else:
# Use extra_params to safely pass query parameters
return self._get_paginated("/images", "images", extra_params={"type": image_type})
def get_available_images(self, image_type: str = "distribution") -> list[dict[str, Any]]:
"""
Get only available distribution images.
Args:
image_type: Filter by image type ('distribution', 'application', or 'all')
Returns:
List of available image objects
"""
images = self.get_images(image_type)
# Filter for public images that are available
return [
img for img in images if img.get("public", False) and img.get("status") == "available"
]
def create_droplet(
self,
name: str,
region: str,
size: str,
image: str,
user_data: str,
tags: list[str],
ssh_keys: list[int] | None = None,
) -> dict[str, Any]:
"""
Create a new droplet.
Args:
name: Droplet name
region: Region slug (e.g., 'nyc3')
size: Size slug (e.g., 's-2vcpu-4gb')
image: Image slug (e.g., 'ubuntu-25-04-x64')
user_data: Cloud-init user data
tags: List of tags to apply
ssh_keys: List of SSH key IDs for root access (optional)
Returns:
Droplet object from API response
"""
payload: dict[str, Any] = {
"name": name,
"region": region,
"size": size,
"image": image,
"user_data": user_data,
"tags": tags,
}
if ssh_keys:
payload["ssh_keys"] = ssh_keys
response = self._request("POST", "/droplets", json=payload)
return response.get("droplet", {})
def get_droplet(self, droplet_id: int) -> dict[str, Any]:
"""
Get droplet information by ID.
Args:
droplet_id: Droplet ID
Returns:
Droplet object
Raises:
ValueError: If droplet_id is not positive
"""
self._validate_positive_int(droplet_id, "droplet_id")
response = self._request("GET", f"/droplets/{droplet_id}")
return response.get("droplet", {})
def list_droplets(self, tag_name: str | None = None) -> list[dict[str, Any]]:
"""
List all droplets, optionally filtered by tag.
Args:
tag_name: Optional tag to filter by (e.g., 'owner:myname')
Returns:
List of droplet objects
"""
if tag_name:
# Use tag-based filtering with safe parameter encoding
return self._get_paginated("/droplets", "droplets", extra_params={"tag_name": tag_name})
else:
# Get all droplets
return self._get_paginated("/droplets", "droplets")
def wait_for_droplet_active(
self,
droplet_id: int,
timeout: int = 300,
poll_interval: int = 5,
) -> dict[str, Any]:
"""
Wait for droplet to become active.
Args:
droplet_id: Droplet ID
timeout: Maximum time to wait in seconds (default 300)
poll_interval: Time between polls in seconds (default 5)
Returns:
Final droplet object
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If timeout is reached or droplet enters error state
"""
import time
self._validate_positive_int(droplet_id, "droplet_id")
start_time = time.time()
while True:
droplet = self.get_droplet(droplet_id)
status = droplet.get("status", "")
if status == "active":
return droplet
elif status == "error":
raise DigitalOceanAPIError(
f"Droplet entered error state: {droplet.get('name', droplet_id)}"
)
elapsed = time.time() - start_time
if elapsed > timeout:
raise DigitalOceanAPIError(
f"Timeout waiting for droplet to become active (waited {elapsed:.0f}s)"
)
time.sleep(poll_interval)
def list_ssh_keys(self) -> list[dict[str, Any]]:
"""
List all SSH keys in the account.
Returns:
List of SSH key objects with id, fingerprint, public_key, name
"""
return self._get_paginated("/account/keys", "ssh_keys")
def get_ssh_key_by_fingerprint(self, fingerprint: str) -> dict[str, Any] | None:
"""
Get SSH key by fingerprint.
Args:
fingerprint: SSH key fingerprint (MD5 format: aa:bb:cc:...)
Returns:
SSH key object if found, None if not found
"""
try:
response = self._request("GET", f"/account/keys/{fingerprint}")
return response.get("ssh_key", {})
except DigitalOceanAPIError as e:
# 404 means key doesn't exist
if e.status_code == 404:
return None
raise
def add_ssh_key(self, name: str, public_key: str) -> dict[str, Any]:
"""
Add a new SSH key to the account.
Args:
name: Name for the SSH key
public_key: The full SSH public key content
Returns:
SSH key object from API response
"""
payload = {
"name": name,
"public_key": public_key,
}
response = self._request("POST", "/account/keys", json=payload)
return response.get("ssh_key", {})
def update_ssh_key(self, key_id: int, name: str) -> dict[str, Any]:
"""
Update an SSH key name.
Args:
key_id: The SSH key ID to update
name: New name for the SSH key
Returns:
Updated SSH key object from API response
Raises:
ValueError: If key_id is not positive
DigitalOceanAPIError: If update fails
"""
self._validate_positive_int(key_id, "key_id")
payload = {"name": name}
response = self._request("PUT", f"/account/keys/{key_id}", json=payload)
return response.get("ssh_key", {})
def delete_ssh_key(self, key_id: int) -> None:
"""
Delete an SSH key from the account.
Args:
key_id: The SSH key ID to delete
Raises:
ValueError: If key_id is not positive
DigitalOceanAPIError: If deletion fails
"""
self._validate_positive_int(key_id, "key_id")
self._request("DELETE", f"/account/keys/{key_id}")
def delete_droplet(self, droplet_id: int) -> None:
"""
Delete a droplet by ID.
Args:
droplet_id: Droplet ID to delete
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If deletion fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
self._request("DELETE", f"/droplets/{droplet_id}")
def rename_droplet(self, droplet_id: int, new_name: str) -> dict[str, Any]:
"""
Rename a droplet.
Args:
droplet_id: Droplet ID
new_name: New name for the droplet
Returns:
Action object with id, status, etc.
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If rename fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
payload = {"type": "rename", "name": new_name}
response = self._request("POST", f"/droplets/{droplet_id}/actions", json=payload)
return response.get("action", {})
def resize_droplet(self, droplet_id: int, size: str, disk: bool = True) -> dict[str, Any]:
"""
Resize a droplet (requires power off, causes downtime).
Args:
droplet_id: Droplet ID
size: New size slug (e.g., 's-4vcpu-8gb')
disk: Whether to resize disk (permanent, cannot be undone). Default: True
Returns:
Action object with id, status, etc.
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If resize fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
payload = {
"type": "resize",
"size": size,
"disk": disk,
}
response = self._request("POST", f"/droplets/{droplet_id}/actions", json=payload)
return response.get("action", {})
def power_on_droplet(self, droplet_id: int) -> dict[str, Any]:
"""
Power on a droplet.
Args:
droplet_id: Droplet ID
Returns:
Action object with id, status, etc.
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If power on fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
payload = {"type": "power_on"}
response = self._request("POST", f"/droplets/{droplet_id}/actions", json=payload)
return response.get("action", {})
def power_off_droplet(self, droplet_id: int) -> dict[str, Any]:
"""
Power off a droplet.
Args:
droplet_id: Droplet ID
Returns:
Action object with id, status, etc.
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If power off fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
payload = {"type": "power_off"}
response = self._request("POST", f"/droplets/{droplet_id}/actions", json=payload)
return response.get("action", {})
def get_action(self, action_id: int) -> dict[str, Any]:
"""
Get action status by ID.
Args:
action_id: Action ID
Returns:
Action object with id, status, type, etc.
Raises:
ValueError: If action_id is not positive
DigitalOceanAPIError: If request fails
"""
self._validate_positive_int(action_id, "action_id")
response = self._request("GET", f"/actions/{action_id}")
return response.get("action", {})
def list_droplet_actions(self, droplet_id: int) -> list[dict[str, Any]]:
"""
List all actions for a droplet.
Args:
droplet_id: Droplet ID
Returns:
List of action objects (most recent first)
Raises:
ValueError: If droplet_id is not positive
"""
self._validate_positive_int(droplet_id, "droplet_id")
return self._get_paginated(f"/droplets/{droplet_id}/actions", "actions")
def wait_for_action_complete(
self,
action_id: int,
timeout: int = 300,
poll_interval: int = 5,
) -> dict[str, Any]:
"""
Wait for action to complete.
Args:
action_id: Action ID
timeout: Maximum time to wait in seconds (default 300)
poll_interval: Time between polls in seconds (default 5)
Returns:
Final action object
Raises:
ValueError: If action_id is not positive
DigitalOceanAPIError: If timeout is reached or action enters error state
"""
import time
self._validate_positive_int(action_id, "action_id")
start_time = time.time()
while True:
action = self.get_action(action_id)
status = action.get("status", "")
if status == "completed":
return action
elif status == "errored":
raise DigitalOceanAPIError(
f"Action failed: {action.get('type', 'unknown')} (ID: {action_id})"
)
elapsed = time.time() - start_time
if elapsed > timeout:
raise DigitalOceanAPIError(
f"Timeout waiting for action to complete (waited {elapsed:.0f}s)"
)
time.sleep(poll_interval)
def list_projects(self) -> list[dict[str, Any]]:
"""
List all projects in the account.
Returns:
List of project objects with id, name, description, purpose, etc.
"""
return self._get_paginated("/projects", "projects")
def get_project(self, project_id: str) -> dict[str, Any] | None:
"""
Get a specific project by ID.
Args:
project_id: Project UUID
Returns:
Project object if found, None if not found (404)
Raises:
DigitalOceanAPIError: If request fails (non-404 errors)
"""
try:
response = self._request("GET", f"/projects/{project_id}")
return response.get("project", {})
except DigitalOceanAPIError as e:
# 404 means project doesn't exist
if e.status_code == 404:
return None
raise
def get_default_project(self) -> dict[str, Any] | None:
"""
Get the default project for the account.
Returns:
Project object if default project exists, None if not found
Raises:
DigitalOceanAPIError: If request fails (non-404 errors)
"""
try:
response = self._request("GET", "/projects/default")
return response.get("project", {})
except DigitalOceanAPIError as e:
# 404 means no default project
if e.status_code == 404:
return None
raise
def assign_resources_to_project(
self, project_id: str, resource_urns: list[str]
) -> dict[str, Any]:
"""
Assign resources to a project.
Args:
project_id: Project UUID
resource_urns: List of resource URNs (e.g., ["do:droplet:12345"])
Returns:
Response with assigned resources
Raises:
DigitalOceanAPIError: If assignment fails
"""
payload = {"resources": resource_urns}
response = self._request("POST", f"/projects/{project_id}/resources", json=payload)
return response
@staticmethod
def get_droplet_urn(droplet_id: int) -> str:
"""
Get the URN (Uniform Resource Name) for a droplet.
Args:
droplet_id: Droplet ID
Returns:
URN string in format "do:droplet:{id}"
"""
return f"do:droplet:{droplet_id}"
# Snapshot methods
def create_snapshot(self, droplet_id: int, name: str) -> dict[str, Any]:
"""
Create a snapshot of a droplet.
Args:
droplet_id: Droplet ID to snapshot
name: Name for the snapshot
Returns:
Action object with id, status, etc.
Raises:
ValueError: If droplet_id is not positive
DigitalOceanAPIError: If snapshot creation fails
"""
self._validate_positive_int(droplet_id, "droplet_id")
payload = {
"type": "snapshot",
"name": name,
}
response = self._request("POST", f"/droplets/{droplet_id}/actions", json=payload)
return response.get("action", {})
def list_snapshots(self, tag: str | None = None) -> list[dict[str, Any]]:
"""
List snapshots, optionally filtered by tag.
Args:
tag: Optional tag to filter by (e.g., 'owner:myname')
Returns:
List of snapshot objects
"""
extra_params: dict[str, str] = {"resource_type": "droplet"}
if tag:
extra_params["tag_name"] = tag
return self._get_paginated("/snapshots", "snapshots", extra_params=extra_params)
def get_snapshot(self, snapshot_id: int) -> dict[str, Any] | None:
"""
Get a snapshot by ID.
Args:
snapshot_id: Snapshot ID
Returns:
Snapshot object if found, None if not found (404)
Raises:
ValueError: If snapshot_id is not positive
DigitalOceanAPIError: If request fails (non-404 errors)
"""
self._validate_positive_int(snapshot_id, "snapshot_id")
try:
response = self._request("GET", f"/snapshots/{snapshot_id}")
return response.get("snapshot", {})
except DigitalOceanAPIError as e:
if e.status_code == 404:
return None
raise
def get_snapshot_by_name(self, name: str, tag: str | None = None) -> dict[str, Any] | None:
"""
Find a snapshot by exact name.
Args:
name: Exact snapshot name to search for
tag: Optional tag to filter by
Returns:
Snapshot object if found, None if not found
"""
snapshots = self.list_snapshots(tag=tag)
for snapshot in snapshots:
if snapshot.get("name") == name:
return snapshot
return None
def delete_snapshot(self, snapshot_id: int) -> None:
"""
Delete a snapshot by ID.
Args:
snapshot_id: Snapshot ID to delete
Raises:
ValueError: If snapshot_id is not positive
DigitalOceanAPIError: If deletion fails
"""
self._validate_positive_int(snapshot_id, "snapshot_id")
self._request("DELETE", f"/snapshots/{snapshot_id}")
def tag_resource(self, tag_name: str, resource_id: str, resource_type: str) -> None:
"""
Add a tag to a resource.
Args:
tag_name: Tag name to apply
resource_id: Resource ID (string for snapshots/images)
resource_type: Resource type ('image' for snapshots, 'droplet', etc.)
Raises:
DigitalOceanAPIError: If tagging fails
"""
payload = {
"resources": [
{
"resource_id": resource_id,
"resource_type": resource_type,
}
]
}
self._request("POST", f"/tags/{tag_name}/resources", json=payload)
def create_tag(self, tag_name: str) -> dict[str, Any]:
"""
Create a tag if it doesn't exist.
Args:
tag_name: Tag name to create
Returns:
Tag object from API response
Raises:
DigitalOceanAPIError: If creation fails (ignores 422 for existing tags)
"""
payload = {"name": tag_name}
try:
response = self._request("POST", "/tags", json=payload)
return response.get("tag", {})
except DigitalOceanAPIError as e:
# 422 means tag already exists, which is fine
if e.status_code == 422:
return {"name": tag_name}
raise
def create_droplet_from_snapshot(
self,
name: str,
region: str,
size: str,
snapshot_id: int,
tags: list[str],
ssh_keys: list[int] | None = None,
) -> dict[str, Any]:
"""
Create a new droplet from a snapshot image.
Args:
name: Droplet name
region: Region slug (e.g., 'nyc3')
size: Size slug (e.g., 's-2vcpu-4gb')
snapshot_id: Snapshot ID to restore from
tags: List of tags to apply
ssh_keys: List of SSH key IDs for root access (optional)
Returns:
Droplet object from API response
Raises:
ValueError: If snapshot_id is not positive
DigitalOceanAPIError: If droplet creation fails
"""
self._validate_positive_int(snapshot_id, "snapshot_id")
payload: dict[str, Any] = {
"name": name,
"region": region,
"size": size,
"image": snapshot_id, # Snapshot ID as the image
"tags": tags,
}
if ssh_keys:
payload["ssh_keys"] = ssh_keys
response = self._request("POST", "/droplets", json=payload)
return response.get("droplet", {})
dropkit/cloudinit.py
+59
View File
@@ -0,0 +1,59 @@
"""Cloud-init template rendering."""
from pathlib import Path
from jinja2 import Template
from dropkit.config import Config
def render_cloud_init(
template_path: str,
username: str,
full_name: str,
email: str,
ssh_keys: list[str],
tailscale_enabled: bool = True,
) -> str:
"""
Render cloud-init template with user data.
Args:
template_path: Path to the cloud-init template file
username: Username to create in the droplet
full_name: Full name extracted from email (for git user.name)
email: Email address from DigitalOcean account (for git user.email)
ssh_keys: List of SSH public key file paths
tailscale_enabled: Whether to install Tailscale VPN (default: True)
Returns:
Rendered cloud-init configuration as string
"""
# Read template
template_file = Path(template_path).expanduser()
if not template_file.exists():
raise FileNotFoundError(f"Cloud-init template not found: {template_path}")
with open(template_file) as f:
template_content = f.read()
# Validate and read SSH key contents
ssh_key_contents = []
for key_path in ssh_keys:
# Validate it's a public key
Config.validate_ssh_public_key(key_path)
# Read the content
content = Config.read_ssh_key_content(key_path)
ssh_key_contents.append(content)
# Render template with Jinja2
template = Template(template_content)
rendered = template.render(
username=username,
full_name=full_name,
email=email,
ssh_keys=ssh_key_contents,
tailscale_enabled=tailscale_enabled,
)
return rendered
dropkit/config.py
+347
View File
@@ -0,0 +1,347 @@
"""Configuration management for dropkit."""
import base64
import hashlib
import os
from pathlib import Path
import yaml
from cryptography.hazmat.primitives import serialization
from pydantic import BaseModel, ConfigDict, Field, field_validator
class DigitalOceanConfig(BaseModel):
"""DigitalOcean API configuration."""
token: str = Field(..., min_length=1, description="DigitalOcean API token")
api_base: str = Field(default="https://api.digitalocean.com/v2", description="API base URL")
@field_validator("token")
@classmethod
def token_not_empty(cls, v: str) -> str:
"""Validate token is not empty or whitespace."""
if not v or not v.strip():
raise ValueError("Token cannot be empty")
return v.strip()
class DefaultsConfig(BaseModel):
"""Default settings for droplet creation."""
region: str = Field(..., min_length=1, description="Default region slug")
size: str = Field(..., min_length=1, description="Default droplet size slug")
image: str = Field(..., min_length=1, description="Default image slug")
extra_tags: list[str] = Field(
default_factory=list,
description="Extra tags (in addition to mandatory owner:<username> and firewall tags)",
)
project_id: str | None = Field(
default=None,
description="Default project ID (UUID) for new droplets",
)
class CloudInitConfig(BaseModel):
"""Cloud-init configuration."""
template_path: str | None = Field(
default=None, description="Path to cloud-init template (None = use package default)"
)
ssh_keys: list[str] = Field(..., min_length=1, description="SSH public key paths")
ssh_key_ids: list[int] = Field(
..., min_length=1, description="DigitalOcean SSH key IDs for root access"
)
@field_validator("ssh_keys")
@classmethod
def ssh_keys_not_empty(cls, v: list[str]) -> list[str]:
"""Validate at least one SSH key is provided."""
if not v:
raise ValueError("At least one SSH key must be configured")
return v
@field_validator("ssh_key_ids")
@classmethod
def ssh_key_ids_not_empty(cls, v: list[int]) -> list[int]:
"""Validate at least one SSH key ID is provided."""
if not v:
raise ValueError("At least one SSH key ID must be configured")
return v
class SSHConfig(BaseModel):
"""SSH configuration."""
config_path: str = Field(..., description="Path to SSH config file")
auto_update: bool = Field(default=True, description="Auto-update SSH config")
identity_file: str = Field(..., description="SSH identity file path")
class TailscaleConfig(BaseModel):
"""Tailscale VPN configuration."""
enabled: bool = Field(default=True, description="Enable Tailscale by default for new droplets")
lock_down_firewall: bool = Field(
default=True, description="Reset UFW to only allow traffic on tailscale0 interface"
)
auth_timeout: int = Field(
default=300, ge=30, description="Timeout in seconds for Tailscale authentication (min: 30)"
)
class DropkitConfig(BaseModel):
"""Main dropkit configuration."""
model_config = ConfigDict(extra="forbid")
digitalocean: DigitalOceanConfig
defaults: DefaultsConfig
cloudinit: CloudInitConfig
ssh: SSHConfig
tailscale: TailscaleConfig = Field(default_factory=TailscaleConfig)
class Config:
"""Manages dropkit configuration with Pydantic validation."""
CONFIG_DIR = Path.home() / ".config" / "dropkit"
CONFIG_FILE = CONFIG_DIR / "config.yaml"
CLOUD_INIT_FILE = CONFIG_DIR / "cloud-init.yaml"
def __init__(self):
"""Initialize config manager."""
self._config: DropkitConfig | None = None
@property
def config(self) -> DropkitConfig:
"""Get the validated configuration."""
if self._config is None:
raise ValueError("Configuration not loaded. Call load() first.")
return self._config
@classmethod
def exists(cls) -> bool:
"""Check if config file exists."""
return cls.CONFIG_FILE.exists()
@classmethod
def get_config_dir(cls) -> Path:
"""Get the config directory path."""
return cls.CONFIG_DIR
@staticmethod
def get_default_template_path() -> Path:
"""Get the default cloud-init template path from package."""
return Path(__file__).parent / "templates" / "default-cloud-init.yaml"
@classmethod
def ensure_config_dir(cls) -> None:
"""Create config directory if it doesn't exist."""
cls.CONFIG_DIR.mkdir(parents=True, exist_ok=True)
# Set restrictive permissions on config directory
cls.CONFIG_DIR.chmod(0o700)
def load(self) -> None:
"""Load and validate configuration from file."""
if not self.CONFIG_FILE.exists():
raise FileNotFoundError(
f"Config file not found at {self.CONFIG_FILE}. Run 'dropkit init' first."
)
with open(self.CONFIG_FILE) as f:
data = yaml.safe_load(f) or {}
# Validate with Pydantic
self._config = DropkitConfig(**data)
def save(self) -> None:
"""Save configuration to file."""
if self._config is None:
raise ValueError("No configuration to save")
self.ensure_config_dir()
# Convert Pydantic model to dict for YAML serialization
data = self._config.model_dump(mode="python")
with open(self.CONFIG_FILE, "w") as f:
yaml.dump(data, f, default_flow_style=False, sort_keys=False)
# Set restrictive permissions on config file (contains API token)
self.CONFIG_FILE.chmod(0o600)
@staticmethod
def get_system_username() -> str:
"""
Get current system username from environment.
This is used for tagging droplets (owner:<username>) to track
who created them.
Returns:
Current system username from $USER environment variable
"""
return os.environ.get("USER", "user")
@staticmethod
def detect_ssh_keys() -> list[str]:
"""Auto-detect all SSH public keys (*.pub) in ~/.ssh/."""
ssh_dir = Path.home() / ".ssh"
if not ssh_dir.exists():
return []
# Find all .pub files in the SSH directory
found_keys = [str(key_path) for key_path in ssh_dir.glob("*.pub")]
# Sort by modification time (most recently modified first)
found_keys.sort(key=lambda p: Path(p).stat().st_mtime, reverse=True)
return found_keys
@staticmethod
def validate_ssh_public_key(key_path: str) -> None:
"""
Validate that a file is a valid SSH public key.
Args:
key_path: Path to the SSH key file
Raises:
ValueError: If the file is not a valid SSH public key
FileNotFoundError: If the file doesn't exist
"""
path = Path(key_path).expanduser()
if not path.exists():
raise FileNotFoundError(f"SSH key not found: {key_path}")
# Check if filename ends with .pub
if not path.name.endswith(".pub"):
raise ValueError(
f"SSH key file must be a public key (*.pub): {key_path}\n"
f"Private keys should NEVER be uploaded to DigitalOcean."
)
# Read and validate content
try:
with open(path) as f:
content = f.read().strip()
except Exception as e:
raise ValueError(f"Cannot read SSH key file: {e}")
if not content:
raise ValueError(f"SSH key file is empty: {key_path}")
# Valid public key prefixes
valid_prefixes = (
"ssh-rsa ",
"ssh-ed25519 ",
"ecdsa-sha2-nistp256 ",
"ecdsa-sha2-nistp384 ",
"ecdsa-sha2-nistp521 ",
"sk-ssh-ed25519@openssh.com ",
"sk-ecdsa-sha2-nistp256@openssh.com ",
)
if not content.startswith(valid_prefixes):
raise ValueError(
f"File does not appear to be a valid SSH public key: {key_path}\n"
f"Public keys should start with: {', '.join(p.strip() for p in valid_prefixes)}"
)
@staticmethod
def read_ssh_key_content(key_path: str) -> str:
"""Read SSH public key content."""
path = Path(key_path).expanduser()
if not path.exists():
raise FileNotFoundError(f"SSH key not found: {key_path}")
with open(path) as f:
return f.read().strip()
@staticmethod
def compute_ssh_key_fingerprint(public_key: str) -> str:
"""
Compute MD5 fingerprint of an SSH public key.
Args:
public_key: SSH public key content
Returns:
MD5 fingerprint in format: aa:bb:cc:dd:...
Raises:
ValueError: If key format is invalid
"""
try:
# Validate and load SSH public key using cryptography library
serialization.load_ssh_public_key(public_key.encode())
# Parse SSH public key format: "ssh-rsa AAAAB3... comment"
# Extract the base64 key material (second field)
parts = public_key.strip().split()
if len(parts) < 2:
raise ValueError("Invalid SSH public key format")
# Decode the base64 key material and compute MD5
key_data = base64.b64decode(parts[1])
fingerprint = hashlib.md5(key_data).hexdigest()
# Format as colon-separated hex pairs
return ":".join(fingerprint[i : i + 2] for i in range(0, len(fingerprint), 2))
except Exception as e:
raise ValueError(f"Failed to compute SSH key fingerprint: {e}")
def create_default_config(
self,
token: str,
username: str, # noqa: ARG002 - kept for API compatibility, username derived from DO API
region: str = "nyc3",
size: str = "s-2vcpu-4gb",
image: str = "ubuntu-25-04-x64",
ssh_keys: list[str] | None = None,
ssh_key_ids: list[int] | None = None,
extra_tags: list[str] | None = None,
project_id: str | None = None,
) -> None:
"""Create a default configuration with validation.
Note: The mandatory tags owner:<username> and firewall are NOT stored in config.
They are always added at runtime when creating droplets.
"""
if ssh_keys is None:
ssh_keys = self.detect_ssh_keys()
if not ssh_keys:
raise ValueError("No SSH keys found. Please specify SSH key paths.")
if ssh_key_ids is None:
raise ValueError("SSH key IDs must be provided from DigitalOcean.")
# Store only extra tags (mandatory tags added at runtime)
if extra_tags is None:
extra_tags = []
# Create validated Pydantic model
self._config = DropkitConfig(
digitalocean=DigitalOceanConfig(
token=token,
api_base="https://api.digitalocean.com/v2",
),
defaults=DefaultsConfig(
region=region,
size=size,
image=image,
extra_tags=extra_tags,
project_id=project_id,
),
cloudinit=CloudInitConfig(
ssh_keys=ssh_keys,
ssh_key_ids=ssh_key_ids,
),
ssh=SSHConfig(
config_path=str(Path.home() / ".ssh" / "config"),
auto_update=True,
identity_file=ssh_keys[0].replace(".pub", "") if ssh_keys else "~/.ssh/id_ed25519",
),
)
dropkit/lock.py
+186
View File
@@ -0,0 +1,186 @@
"""File-based locking to prevent concurrent dropkit write operations."""
import fcntl
import json
import os
import time
from collections.abc import Generator
from contextlib import contextmanager, suppress
from functools import wraps
from pathlib import Path
import typer
from rich.console import Console
# Lock file location
LOCK_FILE = Path("/tmp/dropkit.lock")
# Default timeout for acquiring lock (seconds)
DEFAULT_TIMEOUT = 30.0
# Polling interval when waiting for lock (seconds)
POLL_INTERVAL = 0.5
console = Console()
class LockError(Exception):
"""Raised when lock acquisition fails."""
class LockInfo:
"""Information about the current lock holder."""
def __init__(self, pid: int, command: str):
self.pid = pid
self.command = command
@classmethod
def from_file(cls, lock_file: Path) -> "LockInfo | None":
"""Read lock info from file, returns None if unreadable."""
try:
if lock_file.exists():
content = lock_file.read_text().strip()
if content:
data = json.loads(content)
return cls(pid=data.get("pid", 0), command=data.get("command", "unknown"))
except (json.JSONDecodeError, OSError, KeyError):
pass
return None
def to_json(self) -> str:
"""Serialize lock info to JSON."""
return json.dumps({"pid": self.pid, "command": self.command})
def is_process_alive(self) -> bool:
"""Check if the lock holder process is still running."""
try:
os.kill(self.pid, 0) # Signal 0 checks existence without killing
return True
except OSError:
return False
def _write_lock_info(lock_fd: int, command: str) -> None:
"""Write lock holder information to the lock file."""
info = LockInfo(pid=os.getpid(), command=command)
content = info.to_json().encode()
os.ftruncate(lock_fd, 0)
os.lseek(lock_fd, 0, os.SEEK_SET)
os.write(lock_fd, content)
def _clear_lock_info(lock_fd: int) -> None:
"""Clear lock holder information (on release)."""
with suppress(OSError):
os.ftruncate(lock_fd, 0)
@contextmanager
def operation_lock(
command: str,
timeout: float | None = None,
) -> Generator[None, None, None]:
"""
Context manager that acquires an exclusive lock for dropkit operations.
Uses fcntl.flock() which automatically releases the lock when:
- The context manager exits normally
- An exception is raised
- The process crashes or is killed
Args:
command: Name of the command acquiring the lock (for debugging)
timeout: Maximum seconds to wait for lock (default: 30)
Raises:
LockError: If lock cannot be acquired within timeout
Example:
with operation_lock("create"):
# Exclusive operation here
pass
"""
# Resolve default timeout at runtime (allows patching in tests)
if timeout is None:
timeout = DEFAULT_TIMEOUT
# Ensure parent directory exists (should always exist for /tmp)
LOCK_FILE.parent.mkdir(parents=True, exist_ok=True)
# Open file for reading and writing (creates if doesn't exist)
lock_fd = os.open(str(LOCK_FILE), os.O_RDWR | os.O_CREAT, 0o600)
try:
start_time = time.monotonic()
acquired = False
while time.monotonic() - start_time < timeout:
try:
# Try non-blocking exclusive lock
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
acquired = True
break
except OSError:
# Lock is held by another process, wait and retry
time.sleep(POLL_INTERVAL)
if not acquired:
# Timeout - read lock info for error message
lock_info = LockInfo.from_file(LOCK_FILE)
if lock_info and lock_info.is_process_alive():
raise LockError(
f"Another dropkit operation is in progress.\n"
f" Command: {lock_info.command}\n"
f" PID: {lock_info.pid}\n"
f"Wait for it to complete or check if it's stuck."
)
else:
# Stale lock - try one more time (blocking briefly)
try:
fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
acquired = True
except OSError:
raise LockError(
f"Could not acquire lock after {timeout} seconds.\nLock file: {LOCK_FILE}"
)
# Write lock holder info
_write_lock_info(lock_fd, command)
yield # Execute the protected code
finally:
# Clear lock info and release lock
_clear_lock_info(lock_fd)
fcntl.flock(lock_fd, fcntl.LOCK_UN)
os.close(lock_fd)
def requires_lock(command_name: str):
"""
Decorator that wraps a function with operation_lock and handles errors.
Args:
command_name: Name of the command (for lock info)
Example:
@app.command()
@requires_lock("create")
def create(...):
pass
"""
def decorator(func):
@wraps(func)
def wrapper(*args, **kwargs):
try:
with operation_lock(command_name):
return func(*args, **kwargs)
except LockError as e:
console.print(f"[red]Error: {e}[/red]")
raise typer.Exit(1)
return wrapper
return decorator
dropkit/main.py
+4154
View File
File diff suppressed because it is too large Load Diff
dropkit/ssh_config.py
+303
View File
@@ -0,0 +1,303 @@
"""SSH config file management."""
import shutil
from pathlib import Path
def _backup_config(config_file: Path) -> None:
"""
Create a backup of the SSH config file.
Args:
config_file: Path to the SSH config file
"""
if not config_file.exists():
return
backup_file = config_file.parent / f"{config_file.name}.bak"
# Copy file
shutil.copy2(config_file, backup_file)
# Ensure backup has same permissions as original
original_mode = config_file.stat().st_mode
backup_file.chmod(original_mode)
def add_ssh_host(
config_path: str,
host_name: str,
hostname: str,
user: str,
identity_file: str | None = None,
) -> None:
"""
Add or update an SSH host entry in the SSH config file.
Args:
config_path: Path to SSH config file (e.g., ~/.ssh/config)
host_name: Host alias to use (e.g., 'my-droplet')
hostname: IP address or hostname
user: SSH username
identity_file: Path to SSH private key (optional)
"""
config_file = Path(config_path).expanduser()
# Create SSH directory if it doesn't exist
config_file.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
# Backup existing config before modifying
_backup_config(config_file)
# Read existing config
if config_file.exists():
with open(config_file) as f:
existing_content = f.read()
else:
existing_content = ""
# Check if host already exists
if f"Host {host_name}" in existing_content:
# Host exists, update it
lines = existing_content.split("\n")
new_lines = []
skip_until_next_host = False
for line in lines:
if line.startswith("Host "):
if line == f"Host {host_name}":
# Found our host, skip this block
skip_until_next_host = True
else:
# Different host
skip_until_next_host = False
new_lines.append(line)
elif skip_until_next_host:
# Skip lines in the old host block
continue
else:
new_lines.append(line)
existing_content = "\n".join(new_lines).rstrip() + "\n\n"
# Create new host entry
host_entry = f"Host {host_name}\n"
host_entry += f" HostName {hostname}\n"
host_entry += " ForwardAgent yes\n"
host_entry += f" User {user}\n"
if identity_file:
host_entry += f" IdentityFile {identity_file}\n"
# Ensure existing content ends with newline before appending
if existing_content and not existing_content.endswith("\n"):
existing_content += "\n"
# Append new entry
new_content = existing_content + host_entry + "\n"
# Write back to file
with open(config_file, "w") as f:
f.write(new_content)
# Set restrictive permissions
config_file.chmod(0o600)
def get_ssh_host_ip(config_path: str, host_name: str) -> str | None:
"""
Get the HostName (IP address) for an SSH host entry.
Args:
config_path: Path to SSH config file
host_name: Host alias to look up
Returns:
IP address/hostname if found, None otherwise
"""
config_file = Path(config_path).expanduser()
if not config_file.exists():
return None
with open(config_file) as f:
lines = f.readlines()
in_target_host = False
for line in lines:
stripped = line.strip()
# Check for Host directive
if stripped.startswith("Host "):
# Extract host name (handle multiple hosts on same line)
host_part = stripped[5:].strip()
hosts = host_part.split()
in_target_host = host_name in hosts
elif in_target_host and stripped.startswith("HostName "):
# Found HostName in target host block
return stripped[9:].strip()
elif in_target_host and stripped and not stripped.startswith((" ", "\t", "#")):
# Left the host block (non-indented, non-comment line)
# This handles case where there's no HostName
if not line.startswith((" ", "\t")):
in_target_host = False
return None
def host_exists(config_path: str, host_name: str) -> bool:
"""
Check if an SSH host entry exists in the SSH config file.
Args:
config_path: Path to SSH config file
host_name: Host alias to check
Returns:
True if host exists, False otherwise
"""
config_file = Path(config_path).expanduser()
if not config_file.exists():
return False
with open(config_file) as f:
content = f.read()
return f"Host {host_name}" in content
def remove_ssh_host(config_path: str, host_name: str) -> bool:
"""
Remove an SSH host entry from the SSH config file.
Args:
config_path: Path to SSH config file
host_name: Host alias to remove
Returns:
True if host was found and removed, False otherwise
"""
config_file = Path(config_path).expanduser()
if not config_file.exists():
return False
# Backup existing config before modifying
_backup_config(config_file)
# Read existing config
with open(config_file) as f:
lines = f.readlines()
# Find and remove the host block
new_lines = []
skip_until_next_host = False
found = False
for line in lines:
if line.strip().startswith("Host "):
if line.strip() == f"Host {host_name}":
# Found our host, skip this block
skip_until_next_host = True
found = True
else:
# Different host
skip_until_next_host = False
new_lines.append(line)
elif skip_until_next_host:
# Skip lines in the host block (indented lines)
if line.strip() and not line.startswith((" ", "\t")):
# This is not an indented line, so we're past the block
skip_until_next_host = False
new_lines.append(line)
else:
new_lines.append(line)
if found:
# Write back to file
with open(config_file, "w") as f:
f.writelines(new_lines)
return found
def remove_known_hosts_entry(known_hosts_path: str, hostnames: list[str]) -> int:
"""
Remove entries for specified hostnames from known_hosts file.
Args:
known_hosts_path: Path to known_hosts file (e.g., ~/.ssh/known_hosts)
hostnames: List of hostnames/IPs to remove
Returns:
Number of entries removed
"""
known_hosts_file = Path(known_hosts_path).expanduser()
if not known_hosts_file.exists():
return 0
# Backup existing known_hosts before modifying
_backup_config(known_hosts_file)
# Read existing known_hosts
with open(known_hosts_file) as f:
lines = f.readlines()
# Normalize hostnames for matching (lowercase)
hostnames_lower = {h.lower() for h in hostnames}
# Filter out matching entries
new_lines = []
removed_count = 0
for line in lines:
stripped = line.strip()
# Skip empty lines and comments, preserve them
if not stripped or stripped.startswith("#"):
new_lines.append(line)
continue
# Skip hashed entries (can't match them)
if stripped.startswith("|1|"):
new_lines.append(line)
continue
# Parse the first field (comma-separated hostnames)
# Format: hostname[,hostname2,...] keytype key [comment]
parts = stripped.split(None, 1)
if not parts:
new_lines.append(line)
continue
host_field = parts[0]
entry_hosts = host_field.split(",")
# Check if any of our hostnames match any entry host
should_remove = False
for entry_host in entry_hosts:
# Handle bracketed entries like [hostname]:port
check_host = entry_host.lower()
if check_host.startswith("["):
# Extract hostname from [hostname]:port
bracket_end = check_host.find("]")
if bracket_end > 0:
check_host = check_host[1:bracket_end]
if check_host in hostnames_lower:
should_remove = True
break
if should_remove:
removed_count += 1
else:
new_lines.append(line)
# Only write if we removed something
if removed_count > 0:
with open(known_hosts_file, "w") as f:
f.writelines(new_lines)
return removed_count
dropkit/templates/default-cloud-init.yaml
+135
View File
@@ -0,0 +1,135 @@
#cloud-config
# Update and upgrade system packages
package_update: true
package_upgrade: true
disable_root_opts: no-port-forwarding,no-agent-forwarding,no-X11-forwarding
apt:
sources:
docker:
keyid: 9DC858229FC7DD38854AE2D88D81803C0EBFCD88
keyserver: https://download.docker.com/linux/ubuntu/gpg
source: deb [arch=amd64 signed-by=$KEY_FILE] https://download.docker.com/linux/ubuntu $RELEASE stable
# Install required packages
packages:
- ufw
- zsh
- git
- curl
- wget
- build-essential
- fontconfig
- neovim
- apt-transport-https
- ca-certificates
- gnupg
- docker-ce
- docker-ce-cli
- containerd.io
- docker-buildx-plugin
- docker-compose-plugin
users:
- name: {{ username }}
gecos: "{{ username }} user"
shell: /bin/bash
groups: sudo,admin,wheel,docker
sudo: "ALL=(ALL) NOPASSWD:ALL"
ssh_authorized_keys:{% for key in ssh_keys %}
- {{ key }}
{% endfor %}
# Write custom zshrc file
write_files:
- path: /home/{{ username }}/.zshrc
owner: {{ username }}:{{ username }}
permissions: '0644'
defer: true
content: |
# Enable Powerlevel10k instant prompt. Should stay close to the top of ~/.zshrc.
# Initialization code that may require console input (password prompts, [y/n]
# confirmations, etc.) must go above this block; everything else may go below.
if [[ -r "${XDG_CACHE_HOME:-$HOME/.cache}/p10k-instant-prompt-${(%):-%n}.zsh" ]]; then
source "${XDG_CACHE_HOME:-$HOME/.cache}/p10k-instant-prompt-${(%):-%n}.zsh"
fi
DEFAULT_USER="{{ username }}"
# Zprezto + prompt configuration
source "${ZDOTDIR:-$HOME}/.zprezto/init.zsh"
autoload -Uz promptinit
promptinit
prompt powerlevel10k
# You may need to manually set your language environment
export LANG=en_US.UTF-8
export LC_ALL=en_US.UTF-8
# Preferred editor for local and remote sessions
export EDITOR='nvim'
export LESS='-R'
# Disable open-interpreter telemetry
export DISABLE_TELEMETRY=true
alias vim="nvim"
export PATH="$PATH:/home/{{ username }}/.local/bin"
# To customize prompt, run `p10k configure` or edit ~/.p10k.zsh.
[[ ! -f ~/.p10k.zsh ]] || source ~/.p10k.zsh
# Run commands after boot
runcmd:
# Configure UFW firewall
- ufw default deny incoming
- ufw default allow outgoing
- ufw allow OpenSSH
- ufw limit ssh
- ufw --force enable
{% if tailscale_enabled %}
# Install Tailscale via official script (daemon only, authentication handled by dropkit)
- curl -fsSL https://tailscale.com/install.sh | sh
{% endif %}
# Wait for user to be fully created
- |
timeout=30
count=0
while ! id {{ username }} >/dev/null 2>&1; do
if [ $count -ge $timeout ]; then
echo "Timeout waiting for user {{ username }} to be created"
exit 1
fi
sleep 1
count=$((count + 1))
done
echo "User {{ username }} is ready"
# Ensure home directory exists and has correct permissions
- mkdir -p /home/{{ username }}
- chown {{ username }}:{{ username }} /home/{{ username }}
# Configure git globally for the user
- su {{ username }} -c "git config --global user.name '{{ full_name }}'"
- su {{ username }} -c "git config --global user.email '{{ email }}'"
# Install zprezto for user
- su {{ username }} -c "zsh -c 'git clone --recursive https://github.com/sorin-ionescu/prezto.git /home/{{ username }}/.zprezto'"
# Create Zsh configuration using setopt EXTENDED_GLOB and loop through runcoms
- |
su {{ username }} -c "zsh -c '
setopt EXTENDED_GLOB
for rcfile in /home/{{ username }}/.zprezto/runcoms/^README.md(.N); do
ln -s "\$rcfile" "/home/{{ username }}/.\${rcfile:t}"
done
'"
# Set Zsh as default shell and ensure ownership
- chsh -s /usr/bin/zsh {{ username }}
- chown -R {{ username }}:{{ username }} /home/{{ username }}
- reboot
dropkit/ui.py
+143
View File
@@ -0,0 +1,143 @@
"""UI utilities for dropkit - display functions and prompts."""
from collections.abc import Callable
from typing import Any
from rich.console import Console
from rich.prompt import Prompt
from rich.table import Table
console = Console()
def display_regions(regions: list[dict[str, Any]]) -> None:
"""Display available regions in a table, sorted alphabetically by slug."""
table = Table(title="Available Regions", show_header=True)
table.add_column("Slug", style="cyan", no_wrap=True)
table.add_column("Name", style="white")
table.add_column("Features", style="dim")
# Sort regions alphabetically by slug
sorted_regions = sorted(regions, key=lambda r: r.get("slug", ""))
for region in sorted_regions:
slug = region.get("slug", "")
name = region.get("name", "")
features = ", ".join(region.get("features", [])[:3]) # Show first 3 features
if len(region.get("features", [])) > 3:
features += "..."
table.add_row(slug, name, features)
console.print(table)
def display_sizes(sizes: list[dict[str, Any]]) -> None:
"""Display available droplet sizes in a table, sorted by price."""
table = Table(title="Available Droplet Sizes", show_header=True)
table.add_column("Slug", style="cyan", no_wrap=True)
table.add_column("Memory", style="white", justify="right")
table.add_column("vCPUs", style="white", justify="right")
table.add_column("Disk", style="white", justify="right")
table.add_column("Transfer", style="white", justify="right")
table.add_column("Price/mo", style="green", justify="right")
# Sort sizes by price (monthly) ascending
sorted_sizes = sorted(sizes, key=lambda s: s.get("price_monthly", 0))
for size in sorted_sizes:
slug = size.get("slug", "")
memory = f"{size.get('memory', 0)} MB"
vcpus = str(size.get("vcpus", 0))
disk = f"{size.get('disk', 0)} GB"
transfer = f"{size.get('transfer', 0)} TB"
price = f"${size.get('price_monthly', 0):.2f}"
table.add_row(slug, memory, vcpus, disk, transfer, price)
console.print(table)
def display_images(images: list[dict[str, Any]]) -> None:
"""Display available images in a table, sorted by distribution and name."""
table = Table(title="Available Images", show_header=True)
table.add_column("Slug", style="cyan", no_wrap=True)
table.add_column("Name", style="white")
table.add_column("Distribution", style="dim")
# Sort images by distribution, then by name
sorted_images = sorted(
images, key=lambda img: (img.get("distribution", ""), img.get("name", ""))
)
for image in sorted_images:
slug = image.get("slug", "")
name = image.get("name", "")
distribution = image.get("distribution", "")
# Only show images with slugs (not snapshots)
if slug:
table.add_row(slug, name, distribution)
console.print(table)
def display_projects(projects: list[dict[str, Any]]) -> None:
"""Display available projects in a table, sorted alphabetically by name."""
table = Table(title="Available Projects", show_header=True)
table.add_column("ID", style="dim", no_wrap=True)
table.add_column("Name", style="cyan")
table.add_column("Purpose", style="white")
table.add_column("Description", style="dim")
# Sort projects alphabetically by name
sorted_projects = sorted(projects, key=lambda p: p.get("name", "").lower())
for project in sorted_projects:
project_id = project.get("id", "")
name = project.get("name", "")
purpose = project.get("purpose", "")
description = project.get("description", "")
# Truncate description if too long
if len(description) > 50:
description = description[:47] + "..."
table.add_row(project_id, name, purpose, description)
console.print(table)
def prompt_with_help(
prompt_text: str,
default: str,
display_func: Callable[[list[dict[str, Any]]], None] | None = None,
data: list[dict[str, Any]] | None = None,
) -> str:
"""
Prompt user for input with optional help via '?'.
Args:
prompt_text: The prompt to display (without the default/? part)
default: Default value
display_func: Function to call when user enters '?'
data: Data to pass to display_func
Returns:
User's input value
"""
while True:
value = Prompt.ask(
f"[cyan]{prompt_text} (? for help)[/cyan]",
default=default,
)
if value == "?":
if display_func and data is not None:
console.print()
display_func(data)
console.print()
else:
console.print("[yellow]No help available[/yellow]")
else:
return value
dropkit/version_check.py
+187
View File
@@ -0,0 +1,187 @@
"""Version checking functionality for dropkit."""
import json
import subprocess
import time
from pathlib import Path
from dropkit import __version__
from dropkit.config import Config
def get_last_check_file() -> Path:
"""Get the path to the last version check file."""
return Config.get_config_dir() / ".last_version_check"
def should_check_version() -> bool:
"""
Check if we should check for a new version.
Only checks once per day to avoid slowing down commands.
"""
check_file = get_last_check_file()
if not check_file.exists():
return True
try:
with open(check_file) as f:
data = json.load(f)
last_check = data.get("timestamp", 0)
# Check if more than 24 hours have passed
return (time.time() - last_check) > 86400 # 24 hours in seconds
except (json.JSONDecodeError, OSError):
return True
def update_last_check_time() -> None:
"""Update the last version check timestamp."""
check_file = get_last_check_file()
check_file.parent.mkdir(parents=True, exist_ok=True)
data = {"timestamp": time.time(), "current_version": __version__}
try:
with open(check_file, "w") as f:
json.dump(data, f)
except OSError:
# Silently fail if we can't write the file
pass
def get_latest_git_commit() -> str | None:
"""
Get the latest git commit hash from the main branch.
Returns:
Latest commit hash (short, 7 chars) or None if unable to fetch
"""
try:
# Get the latest commit from main branch
result = subprocess.run(
[
"git",
"ls-remote",
"https://github.com/trailofbits/dropkit.git",
"HEAD",
],
capture_output=True,
timeout=5,
text=True,
)
if result.returncode != 0:
return None
# Parse the output to get the commit hash
# Format: "commit_hash\tHEAD"
lines = result.stdout.strip().split("\n")
if not lines or not lines[0]:
return None
# Extract commit hash (first column)
commit_hash = lines[0].split("\t")[0]
# Return short hash (first 7 chars)
return commit_hash[:7] if commit_hash else None
except (subprocess.TimeoutExpired, subprocess.SubprocessError, OSError):
return None
def extract_commit_from_version(version: str) -> str | None:
"""
Extract commit hash from version string.
Version format: "0.1.0+git.a1b2c3d" or similar
Args:
version: Version string potentially containing a commit hash
Returns:
Commit hash (short) or None if not found
"""
# Version format from hatchling-vcs: "0.1.0+git.a1b2c3d" or "0.1.0.dev123+a1b2c3d"
if "+git." in version:
# Format: "0.1.0+git.a1b2c3d"
parts = version.split("+git.")
if len(parts) == 2:
return parts[1][:7] # Return first 7 chars
elif "+" in version:
# Format: "0.1.0+a1b2c3d" or similar
parts = version.split("+")
if len(parts) == 2:
commit = parts[1]
# Extract hash if it contains other info
if "." in commit:
commit = commit.split(".")[-1]
return commit[:7]
return None
def commits_differ(current_version: str, latest_commit: str) -> bool:
"""
Check if current version's commit differs from latest commit.
Args:
current_version: Current version string (e.g., "0.1.0+git.a1b2c3d")
latest_commit: Latest commit hash from remote
Returns:
True if commits differ (update available), False otherwise
"""
current_commit = extract_commit_from_version(current_version)
if not current_commit:
# Can't determine current commit, don't show update
return False
# Compare commit hashes (case-insensitive)
return current_commit.lower() != latest_commit.lower()
def check_for_updates() -> None:
"""
Check for updates and display a message if a new version is available.
This function:
- Skips check in development mode (version == "dev")
- Only runs once per day for installed versions
- Fetches the latest git commit from main branch
- Compares with current version's commit
- Shows a non-blocking message if commits differ
"""
# Skip check in development mode
if __version__ == "dev" or "dev" in __version__.lower():
return
# Only check once per day
if not should_check_version():
return
# Update the check time regardless of success
update_last_check_time()
# Try to get latest commit
latest_commit = get_latest_git_commit()
if not latest_commit:
# Silently fail if we can't check
return
# Compare commits
if commits_differ(__version__, latest_commit):
# Import here to avoid circular dependency
from rich.console import Console
console = Console()
current_commit = extract_commit_from_version(__version__)
console.print(
f"\n[yellow]New version available:[/yellow] [cyan]{latest_commit}[/cyan] "
f"[dim](current: {current_commit or __version__})[/dim]"
)
console.print(
"[yellow]Run[/yellow] [cyan]uv tool upgrade dropkit[/cyan] [yellow]to update[/yellow]\n"
)
hatch_build.py
+31
View File
@@ -0,0 +1,31 @@
"""Hatchling build hook to embed git commit at build time."""
import subprocess
from pathlib import Path
from hatchling.builders.hooks.plugin.interface import BuildHookInterface
class CustomBuildHook(BuildHookInterface):
"""Build hook to capture git commit hash."""
def initialize(self, version, build_data):
"""Run before the build starts."""
# Get git commit hash
try:
result = subprocess.run(
["git", "rev-parse", "--short=7", "HEAD"],
capture_output=True,
text=True,
timeout=5,
check=False,
)
if result.returncode == 0:
commit = result.stdout.strip()
if commit:
# Write commit to a file that will be included in the package
version_file = Path(self.root) / "dropkit" / "_version.txt"
version_file.write_text(commit)
print(f"Embedded git commit: {commit}")
except Exception as e:
print(f"Warning: Could not capture git commit: {e}")
pyproject.toml
+175
View File
@@ -0,0 +1,175 @@
[project]
name = "dropkit"
version = "0.1.0"
description = "Manage DigitalOcean droplets for ToB engineers"
readme = "README.md"
requires-python = ">=3.11"
authors = [{name = "Trail of Bits", email = "opensource@trailofbits.com"}]
license = {text = "Apache-2.0"}
keywords = ["digitalocean", "droplet", "cli", "devops", "ssh", "tailscale", "cloud"]
classifiers = [
"Development Status :: 4 - Beta",
"Environment :: Console",
"Intended Audience :: Developers",
"License :: OSI Approved :: Apache Software License",
"Operating System :: MacOS",
"Operating System :: POSIX :: Linux",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Topic :: System :: Systems Administration",
]
dependencies = [
"typer>=0.9.0",
"rich>=13.0.0",
"requests>=2.31.0",
"pyyaml>=6.0.1",
"jinja2>=3.1.0",
"shellingham>=1.5.0",
"pydantic>=2.12.3",
"cryptography>=46.0.3",
]
[project.scripts]
dropkit = "dropkit.main:app"
[project.urls]
Homepage = "https://github.com/trailofbits/dropkit"
Repository = "https://github.com/trailofbits/dropkit"
Issues = "https://github.com/trailofbits/dropkit/issues"
[tool.uv]
package = true
default-groups = ["dev"]
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[tool.hatch.build.hooks.custom]
[tool.hatch.build.targets.wheel]
packages = ["dropkit"]
[tool.hatch.build.targets.wheel.force-include]
"dropkit/_version.txt" = "dropkit/_version.txt"
[tool.hatch.build.targets.sdist]
[tool.hatch.build.targets.sdist.force-include]
"dropkit/_version.txt" = "dropkit/_version.txt"
[dependency-groups]
dev = [
"pytest>=8.4.2",
"pytest-cov",
"ruff>=0.8.0",
"ty",
{include-group = "audit"},
]
audit = ["pip-audit"]
[tool.ruff]
target-version = "py311"
line-length = 100
[tool.ruff.format]
quote-style = "double"
indent-style = "space"
line-ending = "auto"
[tool.ruff.lint]
select = ["ALL"]
ignore = [
# Documentation
"D", # pydocstyle - skip for existing code
# Formatter conflicts
"COM812", # missing trailing comma
"ISC001", # implicit string concat
# Line length (handled by formatter)
"E501",
# Exception handling style (CLI prefers clean messages)
"B904", # raise ... from err
"TRY003", # avoid specifying long messages outside exception class
"EM101", # exception must not use string literal
"EM102", # exception must not use f-string literal
"BLE001", # blind exception catching (acceptable in CLI error handling)
# Type annotations (existing code not fully annotated)
"ANN", # flake8-annotations
# Boolean arguments (common in CLI tools)
"FBT001", # boolean positional arg in function definition
"FBT002", # boolean default value in function definition
"FBT003", # boolean positional value in function call
# Magic values (acceptable in existing code)
"PLR2004",
# Import organization (acceptable for lazy imports in CLI)
"PLC0415",
# Complexity (existing code, would require significant refactoring)
"PLR0912", # too many branches
"PLR0913", # too many arguments
"PLR0915", # too many statements
"C901", # function too complex
# Subprocess (intentional CLI tool usage)
"S603", # subprocess call without shell=True check
"S607", # partial executable path
# pytest style
"PT011", # pytest.raises too broad
# Other
"TRY300", # try-except-else instead of try-except
"TRY301", # abstract raise to inner function
"RET504", # unnecessary assignment before return
"RET505", # unnecessary else after return
"RET506", # unnecessary elif after raise
"PLW1510", # subprocess.run without check
"ERA001", # commented out code (false positives on format comments)
"RUF059", # unused unpacked variable (prefixing with _ breaks readability)
"S110", # try-except-pass (acceptable in version detection)
"S108", # hardcoded temp file (intentional for lockfile)
"S324", # insecure hash md5 (used for SSH key fingerprints, industry standard)
"PTH123", # use Path.open instead of open (existing code pattern)
"PGH003", # use specific rule codes (existing code)
"PIE807", # prefer list over lambda (existing code)
]
[tool.ruff.lint.per-file-ignores]
"tests/*" = [
"S101", # assert is fine in tests
"S105", # hardcoded password in tests
"S106", # hardcoded password argument in tests
"ARG", # unused arguments in test fixtures
"SLF001", # private member access in tests
"PLR0913", # too many arguments in test functions
"PT022", # yield vs return in fixtures
]
"hatch_build.py" = [
"T201", # print is fine in build hooks
"ARG002", # unused arguments in hatch hooks
]
[tool.ruff.lint.pyupgrade]
# Enforce Python 3.10+ syntax
keep-runtime-typing = false
[tool.ty.terminal]
error-on-warning = true
[tool.ty.environment]
python-version = "3.11"
[tool.pytest.ini_options]
testpaths = ["tests"]
python_files = ["test_*.py"]
python_classes = ["Test*"]
python_functions = ["test_*"]
addopts = "--cov=dropkit --cov-fail-under=29"
tests/__init__.py
+1
View File
@@ -0,0 +1 @@
"""Tests for dropkit."""
tests/test_api.py
+107
View File
@@ -0,0 +1,107 @@
"""Tests for DigitalOcean API client."""
import pytest
from dropkit.api import DigitalOceanAPI
class TestDigitalOceanAPI:
"""Tests for DigitalOceanAPI class."""
def test_sanitize_email_for_username_trailofbits(self):
"""Test sanitizing trailofbits.com email."""
username = DigitalOceanAPI._sanitize_email_for_username("john.doe@trailofbits.com")
assert username == "john_doe"
def test_sanitize_email_for_username_other_domain(self):
"""Test sanitizing other domain email."""
username = DigitalOceanAPI._sanitize_email_for_username("john.doe@example.com")
assert username == "john_doe"
def test_sanitize_email_for_username_special_chars(self):
"""Test sanitizing email with special characters."""
username = DigitalOceanAPI._sanitize_email_for_username("john-doe+test@trailofbits.com")
assert username == "john_doe_test"
def test_sanitize_email_for_username_starts_with_number(self):
"""Test sanitizing email that starts with number."""
username = DigitalOceanAPI._sanitize_email_for_username("123user@trailofbits.com")
assert username == "u123user"
def test_sanitize_email_for_username_empty_fallback(self):
"""Test sanitizing empty email."""
username = DigitalOceanAPI._sanitize_email_for_username("@trailofbits.com")
assert username == "user"
class TestValidatePositiveInt:
"""Tests for _validate_positive_int method."""
def test_valid_positive_int(self):
"""Test validation passes for positive integer."""
# Should not raise
DigitalOceanAPI._validate_positive_int(1, "test_id")
DigitalOceanAPI._validate_positive_int(100, "test_id")
DigitalOceanAPI._validate_positive_int(999999, "test_id")
def test_zero_raises_error(self):
"""Test validation fails for zero."""
with pytest.raises(ValueError, match="test_id must be a positive integer"):
DigitalOceanAPI._validate_positive_int(0, "test_id")
def test_negative_raises_error(self):
"""Test validation fails for negative integer."""
with pytest.raises(ValueError, match="droplet_id must be a positive integer"):
DigitalOceanAPI._validate_positive_int(-1, "droplet_id")
with pytest.raises(ValueError, match="action_id must be a positive integer"):
DigitalOceanAPI._validate_positive_int(-100, "action_id")
def test_error_message_includes_value(self):
"""Test error message includes the invalid value."""
with pytest.raises(ValueError, match="got: -5"):
DigitalOceanAPI._validate_positive_int(-5, "snapshot_id")
class TestGetDropletUrn:
"""Tests for get_droplet_urn static method."""
def test_urn_format(self):
"""Test URN is in correct format."""
assert DigitalOceanAPI.get_droplet_urn(12345) == "do:droplet:12345"
def test_urn_with_large_id(self):
"""Test URN with large droplet ID."""
assert DigitalOceanAPI.get_droplet_urn(999999999) == "do:droplet:999999999"
class TestListDropletActions:
"""Tests for list_droplet_actions method validation."""
def test_list_droplet_actions_invalid_id_zero(self):
"""Test that list_droplet_actions raises ValueError for zero droplet_id."""
api = DigitalOceanAPI("fake-token")
with pytest.raises(ValueError, match="droplet_id must be a positive integer"):
api.list_droplet_actions(0)
def test_list_droplet_actions_invalid_id_negative(self):
"""Test that list_droplet_actions raises ValueError for negative droplet_id."""
api = DigitalOceanAPI("fake-token")
with pytest.raises(ValueError, match="droplet_id must be a positive integer"):
api.list_droplet_actions(-1)
class TestRenameDroplet:
"""Tests for rename_droplet method validation."""
def test_rename_droplet_invalid_id_zero(self):
"""Test that rename_droplet raises ValueError for zero droplet_id."""
api = DigitalOceanAPI("fake-token")
with pytest.raises(ValueError, match="droplet_id must be a positive integer"):
api.rename_droplet(0, "new-name")
def test_rename_droplet_invalid_id_negative(self):
"""Test that rename_droplet raises ValueError for negative droplet_id."""
api = DigitalOceanAPI("fake-token")
with pytest.raises(ValueError, match="droplet_id must be a positive integer"):
api.rename_droplet(-1, "new-name")
tests/test_config.py
+523
View File
@@ -0,0 +1,523 @@
"""Comprehensive tests for configuration management with Pydantic validation."""
from pathlib import Path
import pytest
from pydantic import ValidationError
from dropkit.config import (
CloudInitConfig,
Config,
DefaultsConfig,
DigitalOceanConfig,
DropkitConfig,
SSHConfig,
)
@pytest.fixture
def temp_config_dir(tmp_path):
"""Create a temporary config directory for testing."""
config_dir = tmp_path / ".config" / "dropkit"
config_dir.mkdir(parents=True)
yield config_dir
@pytest.fixture
def valid_config_dict():
"""Return a valid configuration dictionary."""
return {
"digitalocean": {
"token": "dop_v1_test_token_12345",
"api_base": "https://api.digitalocean.com/v2",
},
"defaults": {
"region": "nyc3",
"size": "s-2vcpu-4gb",
"image": "ubuntu-25-04-x64",
"extra_tags": ["custom-tag"],
},
"cloudinit": {
"template_path": "/home/user/.config/dropkit/cloud-init.yaml",
"ssh_keys": ["/home/user/.ssh/id_ed25519.pub"],
"ssh_key_ids": [12345678],
},
"ssh": {
"config_path": "/home/user/.ssh/config",
"auto_update": True,
"identity_file": "/home/user/.ssh/id_ed25519",
},
}
class TestDigitalOceanConfig:
"""Tests for DigitalOceanConfig validation."""
def test_valid_config(self):
"""Test valid DigitalOcean config."""
config = DigitalOceanConfig(
token="dop_v1_test_token", api_base="https://api.digitalocean.com/v2"
)
assert config.token == "dop_v1_test_token"
assert config.api_base == "https://api.digitalocean.com/v2"
def test_token_whitespace_stripped(self):
"""Test that token whitespace is stripped."""
config = DigitalOceanConfig(token=" dop_v1_test ")
assert config.token == "dop_v1_test"
def test_empty_token_fails(self):
"""Test that empty token fails validation."""
with pytest.raises(ValidationError) as exc_info:
DigitalOceanConfig(token="")
# Pydantic catches this with min_length before our custom validator
assert "token" in str(exc_info.value).lower()
def test_whitespace_only_token_fails(self):
"""Test that whitespace-only token fails validation."""
with pytest.raises(ValidationError) as exc_info:
DigitalOceanConfig(token=" ")
assert "Token cannot be empty" in str(exc_info.value)
def test_missing_token_fails(self):
"""Test that missing token fails validation."""
with pytest.raises(ValidationError):
DigitalOceanConfig()
def test_default_api_base(self):
"""Test default API base URL."""
config = DigitalOceanConfig(token="test_token")
assert config.api_base == "https://api.digitalocean.com/v2"
class TestDefaultsConfig:
"""Tests for DefaultsConfig validation."""
def test_valid_config(self):
"""Test valid defaults config."""
config = DefaultsConfig(
region="nyc3",
size="s-2vcpu-4gb",
image="ubuntu-25-04-x64",
extra_tags=["tag1", "tag2"],
)
assert config.region == "nyc3"
assert config.size == "s-2vcpu-4gb"
assert config.image == "ubuntu-25-04-x64"
assert config.extra_tags == ["tag1", "tag2"]
def test_empty_region_fails(self):
"""Test that empty region fails validation."""
with pytest.raises(ValidationError):
DefaultsConfig(
region="",
size="s-2vcpu-4gb",
image="ubuntu-25-04-x64",
)
def test_empty_size_fails(self):
"""Test that empty size fails validation."""
with pytest.raises(ValidationError):
DefaultsConfig(
region="nyc3",
size="",
image="ubuntu-25-04-x64",
)
def test_empty_image_fails(self):
"""Test that empty image fails validation."""
with pytest.raises(ValidationError):
DefaultsConfig(
region="nyc3",
size="s-2vcpu-4gb",
image="",
)
def test_empty_extra_tags_allowed(self):
"""Test that empty extra_tags list is allowed."""
config = DefaultsConfig(
region="nyc3",
size="s-2vcpu-4gb",
image="ubuntu-25-04-x64",
extra_tags=[],
)
assert config.extra_tags == []
def test_missing_extra_tags_defaults_to_empty(self):
"""Test that missing extra_tags defaults to empty list."""
config = DefaultsConfig(
region="nyc3",
size="s-2vcpu-4gb",
image="ubuntu-25-04-x64",
)
assert config.extra_tags == []
class TestCloudInitConfig:
"""Tests for CloudInitConfig validation."""
def test_valid_config(self):
"""Test valid cloud-init config."""
config = CloudInitConfig(
template_path="/path/to/template.yaml",
ssh_keys=["/path/to/key1.pub", "/path/to/key2.pub"],
ssh_key_ids=[12345, 67890],
)
assert config.template_path == "/path/to/template.yaml"
assert len(config.ssh_keys) == 2
assert len(config.ssh_key_ids) == 2
def test_empty_ssh_keys_fails(self):
"""Test that empty SSH keys list fails validation."""
with pytest.raises(ValidationError) as exc_info:
CloudInitConfig(
template_path="/path/to/template.yaml",
ssh_keys=[],
ssh_key_ids=[12345],
)
# Pydantic catches this with min_length
assert "ssh_keys" in str(exc_info.value).lower()
def test_missing_ssh_keys_fails(self):
"""Test that missing SSH keys fails validation."""
with pytest.raises(ValidationError):
CloudInitConfig(
template_path="/path/to/template.yaml",
ssh_key_ids=[12345],
)
def test_single_ssh_key_valid(self):
"""Test that single SSH key is valid."""
config = CloudInitConfig(
template_path="/path/to/template.yaml",
ssh_keys=["/path/to/key.pub"],
ssh_key_ids=[12345],
)
assert len(config.ssh_keys) == 1
class TestSSHConfig:
"""Tests for SSHConfig validation."""
def test_valid_config(self):
"""Test valid SSH config."""
config = SSHConfig(
config_path="/home/user/.ssh/config",
auto_update=True,
identity_file="/home/user/.ssh/id_ed25519",
)
assert config.config_path == "/home/user/.ssh/config"
assert config.auto_update is True
assert config.identity_file == "/home/user/.ssh/id_ed25519"
def test_auto_update_defaults_to_true(self):
"""Test that auto_update defaults to True."""
config = SSHConfig(
config_path="/home/user/.ssh/config",
identity_file="/home/user/.ssh/id_ed25519",
)
assert config.auto_update is True
def test_missing_config_path_fails(self):
"""Test that missing config_path fails validation."""
with pytest.raises(ValidationError):
SSHConfig(identity_file="/home/user/.ssh/id_ed25519")
def test_missing_identity_file_fails(self):
"""Test that missing identity_file fails validation."""
with pytest.raises(ValidationError):
SSHConfig(config_path="/home/user/.ssh/config")
class TestDropkitConfig:
"""Tests for DropkitConfig (full configuration) validation."""
def test_valid_full_config(self, valid_config_dict):
"""Test valid full configuration."""
config = DropkitConfig(**valid_config_dict)
assert config.digitalocean.token == "dop_v1_test_token_12345"
assert config.defaults.region == "nyc3"
assert config.cloudinit.ssh_keys[0] == "/home/user/.ssh/id_ed25519.pub"
assert config.ssh.auto_update is True
def test_missing_digitalocean_section_fails(self, valid_config_dict):
"""Test that missing digitalocean section fails validation."""
del valid_config_dict["digitalocean"]
with pytest.raises(ValidationError):
DropkitConfig(**valid_config_dict)
def test_missing_defaults_section_fails(self, valid_config_dict):
"""Test that missing defaults section fails validation."""
del valid_config_dict["defaults"]
with pytest.raises(ValidationError):
DropkitConfig(**valid_config_dict)
def test_missing_cloudinit_section_fails(self, valid_config_dict):
"""Test that missing cloudinit section fails validation."""
del valid_config_dict["cloudinit"]
with pytest.raises(ValidationError):
DropkitConfig(**valid_config_dict)
def test_missing_ssh_section_fails(self, valid_config_dict):
"""Test that missing ssh section fails validation."""
del valid_config_dict["ssh"]
with pytest.raises(ValidationError):
DropkitConfig(**valid_config_dict)
def test_extra_fields_forbidden(self, valid_config_dict):
"""Test that extra fields are forbidden."""
valid_config_dict["extra_field"] = "not allowed"
with pytest.raises(ValidationError) as exc_info:
DropkitConfig(**valid_config_dict)
assert "extra_field" in str(exc_info.value).lower()
def test_nested_validation_error(self, valid_config_dict):
"""Test that nested validation errors are caught."""
valid_config_dict["digitalocean"]["token"] = ""
with pytest.raises(ValidationError) as exc_info:
DropkitConfig(**valid_config_dict)
assert "token" in str(exc_info.value).lower()
class TestConfigManager:
"""Tests for Config manager class."""
def test_config_not_loaded_raises_error(self):
"""Test that accessing config before loading raises error."""
config = Config()
with pytest.raises(ValueError) as exc_info:
_ = config.config
assert "not loaded" in str(exc_info.value).lower()
def test_create_default_config(self):
"""Test creating default configuration."""
config = Config()
config.create_default_config(
token="test_token",
username="testuser",
region="nyc3",
size="s-2vcpu-4gb",
image="ubuntu-25-04-x64",
ssh_keys=["/path/to/key.pub"],
ssh_key_ids=[12345],
extra_tags=["tag1", "tag2"],
)
# Should be able to access config now
assert config.config.digitalocean.token == "test_token"
assert config.config.defaults.region == "nyc3"
assert config.config.defaults.extra_tags == ["tag1", "tag2"]
assert config.config.cloudinit.ssh_key_ids == [12345]
def test_create_config_without_ssh_keys_raises_error(self, monkeypatch):
"""Test that creating config without SSH keys raises error."""
# Mock detect_ssh_keys to return empty list
monkeypatch.setattr(Config, "detect_ssh_keys", staticmethod(list))
config = Config()
with pytest.raises(ValueError) as exc_info:
config.create_default_config(
token="test_token",
username="testuser",
ssh_keys=None, # Will try to auto-detect
ssh_key_ids=[12345],
)
assert "SSH key" in str(exc_info.value)
def test_save_without_config_raises_error(self):
"""Test that saving without config raises error."""
config = Config()
with pytest.raises(ValueError) as exc_info:
config.save()
assert "No configuration to save" in str(exc_info.value)
def test_save_and_load_config(self, temp_config_dir, valid_config_dict, monkeypatch):
"""Test saving and loading configuration."""
# Monkey patch Config paths
monkeypatch.setattr(Config, "CONFIG_DIR", temp_config_dir)
monkeypatch.setattr(Config, "CONFIG_FILE", temp_config_dir / "config.yaml")
monkeypatch.setattr(Config, "CLOUD_INIT_FILE", temp_config_dir / "cloud-init.yaml")
# Create and save config
config = Config()
config.create_default_config(
token="test_token_save_load",
username="testuser",
region="sfo3",
size="s-1vcpu-1gb",
image="ubuntu-25-04-x64",
ssh_keys=["/path/to/key.pub"],
ssh_key_ids=[98765],
extra_tags=["test_tag"],
)
config.save()
# Verify file exists
assert Config.CONFIG_FILE.exists()
# Verify file permissions
mode = Config.CONFIG_FILE.stat().st_mode & 0o777
assert mode == 0o600
# Load config in new instance
config2 = Config()
config2.load()
# Verify loaded config matches
assert config2.config.digitalocean.token == "test_token_save_load"
assert config2.config.defaults.region == "sfo3"
assert config2.config.defaults.size == "s-1vcpu-1gb"
assert config2.config.defaults.extra_tags == ["test_tag"]
assert config2.config.cloudinit.ssh_key_ids == [98765]
def test_load_invalid_config_raises_validation_error(self, temp_config_dir, monkeypatch):
"""Test that loading invalid config raises validation error."""
# Monkey patch Config paths
monkeypatch.setattr(Config, "CONFIG_DIR", temp_config_dir)
monkeypatch.setattr(Config, "CONFIG_FILE", temp_config_dir / "config.yaml")
# Write invalid config (missing required fields)
invalid_config = {
"digitalocean": {
"token": "", # Empty token
}
}
import yaml
with open(Config.CONFIG_FILE, "w") as f:
yaml.dump(invalid_config, f)
# Try to load
config = Config()
with pytest.raises(ValidationError):
config.load()
def test_load_nonexistent_config_raises_file_not_found(self, temp_config_dir, monkeypatch):
"""Test that loading non-existent config raises FileNotFoundError."""
# Monkey patch Config paths
monkeypatch.setattr(Config, "CONFIG_DIR", temp_config_dir)
monkeypatch.setattr(Config, "CONFIG_FILE", temp_config_dir / "nonexistent.yaml")
config = Config()
with pytest.raises(FileNotFoundError) as exc_info:
config.load()
assert "dropkit init" in str(exc_info.value)
def test_get_system_username(self, monkeypatch):
"""Test getting system username."""
monkeypatch.setenv("USER", "testuser")
assert Config.get_system_username() == "testuser"
def test_get_system_username_default(self, monkeypatch):
"""Test getting system username with no USER env var."""
monkeypatch.delenv("USER", raising=False)
assert Config.get_system_username() == "user"
def test_detect_ssh_keys(self, tmp_path, monkeypatch):
"""Test SSH key detection."""
# Create fake SSH directory
ssh_dir = tmp_path / ".ssh"
ssh_dir.mkdir()
# Create some keys
(ssh_dir / "id_ed25519.pub").touch()
(ssh_dir / "id_rsa.pub").touch()
# Monkey patch home directory
monkeypatch.setattr(Path, "home", lambda: tmp_path)
keys = Config.detect_ssh_keys()
assert len(keys) == 2
assert str(ssh_dir / "id_ed25519.pub") in keys
assert str(ssh_dir / "id_rsa.pub") in keys
def test_detect_ssh_keys_no_ssh_dir(self, tmp_path, monkeypatch):
"""Test SSH key detection with no .ssh directory."""
# Monkey patch home directory (no .ssh dir)
monkeypatch.setattr(Path, "home", lambda: tmp_path)
keys = Config.detect_ssh_keys()
assert keys == []
def test_read_ssh_key_content(self, tmp_path):
"""Test reading SSH key content."""
key_file = tmp_path / "test_key.pub"
key_file.write_text("ssh-rsa AAAAB3... test@example.com\n")
content = Config.read_ssh_key_content(str(key_file))
assert content == "ssh-rsa AAAAB3... test@example.com"
def test_read_ssh_key_nonexistent_raises_error(self):
"""Test reading non-existent SSH key raises error."""
with pytest.raises(FileNotFoundError):
Config.read_ssh_key_content("/nonexistent/key.pub")
def test_validate_ssh_public_key_valid_ed25519(self, tmp_path):
"""Test validation of valid ED25519 public key."""
key_file = tmp_path / "id_ed25519.pub"
key_file.write_text("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAI... user@example.com\n")
# Should not raise any exception
Config.validate_ssh_public_key(str(key_file))
def test_validate_ssh_public_key_valid_rsa(self, tmp_path):
"""Test validation of valid RSA public key."""
key_file = tmp_path / "id_rsa.pub"
key_file.write_text("ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAAB... user@example.com\n")
# Should not raise any exception
Config.validate_ssh_public_key(str(key_file))
def test_validate_ssh_public_key_valid_ecdsa(self, tmp_path):
"""Test validation of valid ECDSA public key."""
key_file = tmp_path / "id_ecdsa.pub"
key_file.write_text(
"ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNT... user@example.com\n"
)
# Should not raise any exception
Config.validate_ssh_public_key(str(key_file))
def test_validate_ssh_public_key_private_key_fails(self, tmp_path):
"""Test that private key (without .pub extension) fails validation."""
key_file = tmp_path / "id_ed25519"
key_file.write_text("-----BEGIN OPENSSH PRIVATE KEY-----\n...")
with pytest.raises(ValueError) as exc_info:
Config.validate_ssh_public_key(str(key_file))
assert "must be a public key" in str(exc_info.value)
assert "Private keys should NEVER" in str(exc_info.value)
def test_validate_ssh_public_key_nonexistent_fails(self):
"""Test that non-existent key file fails validation."""
with pytest.raises(FileNotFoundError) as exc_info:
Config.validate_ssh_public_key("/nonexistent/key.pub")
assert "not found" in str(exc_info.value)
def test_validate_ssh_public_key_empty_file_fails(self, tmp_path):
"""Test that empty key file fails validation."""
key_file = tmp_path / "empty.pub"
key_file.write_text("")
with pytest.raises(ValueError) as exc_info:
Config.validate_ssh_public_key(str(key_file))
assert "empty" in str(exc_info.value)
def test_validate_ssh_public_key_invalid_content_fails(self, tmp_path):
"""Test that file with invalid content fails validation."""
key_file = tmp_path / "invalid.pub"
key_file.write_text("This is not a valid SSH public key\n")
with pytest.raises(ValueError) as exc_info:
Config.validate_ssh_public_key(str(key_file))
assert "does not appear to be a valid SSH public key" in str(exc_info.value)
def test_validate_ssh_public_key_private_key_header_fails(self, tmp_path):
"""Test that file with private key header fails validation."""
key_file = tmp_path / "private.pub" # Has .pub extension but contains private key
key_file.write_text("-----BEGIN OPENSSH PRIVATE KEY-----\n")
with pytest.raises(ValueError) as exc_info:
Config.validate_ssh_public_key(str(key_file))
assert "does not appear to be a valid SSH public key" in str(exc_info.value)
tests/test_lock.py
+354
View File
@@ -0,0 +1,354 @@
"""Tests for file-based locking mechanism."""
import os
import subprocess
import sys
import time
from pathlib import Path
from unittest.mock import patch
import pytest
from dropkit.lock import (
DEFAULT_TIMEOUT,
LOCK_FILE,
LockError,
LockInfo,
operation_lock,
requires_lock,
)
@pytest.fixture
def temp_lock_file(tmp_path):
"""Use a temporary lock file for testing."""
test_lock = tmp_path / "dropkit.lock"
with patch("dropkit.lock.LOCK_FILE", test_lock):
yield test_lock
class TestLockInfo:
"""Tests for LockInfo class."""
def test_to_json(self):
"""Test serialization to JSON."""
info = LockInfo(pid=12345, command="create")
json_str = info.to_json()
assert '"pid": 12345' in json_str
assert '"command": "create"' in json_str
def test_from_file_valid(self, tmp_path):
"""Test reading valid lock info from file."""
lock_file = tmp_path / "test.lock"
lock_file.write_text('{"pid": 99999, "command": "destroy"}')
info = LockInfo.from_file(lock_file)
assert info is not None
assert info.pid == 99999
assert info.command == "destroy"
def test_from_file_invalid_json(self, tmp_path):
"""Test reading invalid JSON returns None."""
lock_file = tmp_path / "test.lock"
lock_file.write_text("not valid json")
info = LockInfo.from_file(lock_file)
assert info is None
def test_from_file_empty(self, tmp_path):
"""Test reading empty file returns None."""
lock_file = tmp_path / "test.lock"
lock_file.write_text("")
info = LockInfo.from_file(lock_file)
assert info is None
def test_from_file_nonexistent(self, tmp_path):
"""Test reading nonexistent file returns None."""
nonexistent = tmp_path / "does_not_exist.lock"
info = LockInfo.from_file(nonexistent)
assert info is None
def test_is_process_alive_current_process(self):
"""Test that current process is detected as alive."""
info = LockInfo(pid=os.getpid(), command="test")
assert info.is_process_alive() is True
def test_is_process_alive_dead_process(self):
"""Test that non-existent PID is detected as dead."""
# Use a very high PID unlikely to exist
info = LockInfo(pid=9999999, command="test")
assert info.is_process_alive() is False
class TestOperationLock:
"""Tests for operation_lock context manager."""
def test_basic_lock_acquire_release(self, temp_lock_file):
"""Test basic lock acquisition and release."""
with operation_lock("test"):
# Lock should be held
assert temp_lock_file.exists()
content = temp_lock_file.read_text()
assert str(os.getpid()) in content
assert "test" in content
def test_lock_info_cleared_on_exit(self, temp_lock_file):
"""Test lock info is cleared after context exits."""
with operation_lock("test"):
pass
# Lock file should exist but be empty (truncated)
assert temp_lock_file.exists()
content = temp_lock_file.read_text()
assert content == ""
def test_lock_released_on_exception(self, temp_lock_file):
"""Test lock is released when exception is raised."""
with pytest.raises(ValueError), operation_lock("test"):
raise ValueError("test error")
# Lock should be released - can acquire again immediately
with operation_lock("test"):
pass # Should not raise
def test_lock_with_empty_command_name(self, temp_lock_file):
"""Test lock with empty command name."""
with operation_lock(""):
content = temp_lock_file.read_text()
assert '"command": ""' in content
def test_lock_with_special_characters(self, temp_lock_file):
"""Test lock with special characters in command name."""
with operation_lock("enable-tailscale"):
content = temp_lock_file.read_text()
assert "enable-tailscale" in content
class TestConcurrentLock:
"""Tests for concurrent lock behavior using subprocess."""
def test_concurrent_lock_blocks(self, tmp_path):
"""Test that second process waits for first to release lock."""
test_lock = tmp_path / "dropkit.lock"
results_file = tmp_path / "results.txt"
# Script that holds lock and logs events
holder_script = f"""
import sys
import fcntl
import os
import time
lock_path = "{test_lock}"
results_path = "{results_file}"
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
fcntl.flock(fd, fcntl.LOCK_EX)
with open(results_path, "a") as f:
f.write("first_acquired\\n")
f.flush()
time.sleep(1)
with open(results_path, "a") as f:
f.write("first_released\\n")
fcntl.flock(fd, fcntl.LOCK_UN)
os.close(fd)
"""
waiter_script = f"""
import sys
import fcntl
import os
import time
lock_path = "{test_lock}"
results_path = "{results_file}"
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
fcntl.flock(fd, fcntl.LOCK_EX) # Will block until holder releases
with open(results_path, "a") as f:
f.write("second_acquired\\n")
fcntl.flock(fd, fcntl.LOCK_UN)
os.close(fd)
"""
# Start holder first
p1 = subprocess.Popen([sys.executable, "-c", holder_script])
time.sleep(0.3) # Let holder acquire lock
# Start waiter
p2 = subprocess.Popen([sys.executable, "-c", waiter_script])
p1.wait(timeout=10)
p2.wait(timeout=10)
# Read results
events = results_file.read_text().strip().split("\n")
# First should acquire before second
assert "first_acquired" in events
assert "second_acquired" in events
assert events.index("first_acquired") < events.index("second_acquired")
def test_timeout_raises_lock_error(self, tmp_path):
"""Test that timeout raises LockError."""
test_lock = tmp_path / "dropkit.lock"
ready_file = tmp_path / "ready"
# Script that holds lock indefinitely until killed
holder_script = f"""
import fcntl
import os
import time
from pathlib import Path
lock_path = "{test_lock}"
ready_path = "{ready_file}"
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
fcntl.flock(fd, fcntl.LOCK_EX)
Path(ready_path).touch() # Signal we have the lock
time.sleep(30) # Hold lock
"""
p = subprocess.Popen([sys.executable, "-c", holder_script])
try:
# Wait for holder to acquire lock
for _ in range(50):
if ready_file.exists():
break
time.sleep(0.1)
else:
pytest.fail("Holder did not acquire lock in time")
# Now try to acquire with short timeout
with patch("dropkit.lock.LOCK_FILE", test_lock):
with pytest.raises(LockError) as exc_info: # noqa: SIM117
with operation_lock("waiter", timeout=1):
pass
error_msg = str(exc_info.value)
assert (
"Another dropkit operation is in progress" in error_msg
or "Could not acquire lock" in error_msg
)
finally:
p.terminate()
p.wait()
def test_stale_lock_cleanup(self, tmp_path):
"""Test that stale locks from dead processes are handled."""
test_lock = tmp_path / "dropkit.lock"
# Write stale lock info with dead PID
test_lock.write_text('{"pid": 9999999, "command": "stale"}')
# Should still be able to acquire lock
with patch("dropkit.lock.LOCK_FILE", test_lock), operation_lock("new_command"):
content = test_lock.read_text()
assert "new_command" in content
class TestRequiresLockDecorator:
"""Tests for requires_lock decorator."""
def test_decorator_acquires_lock(self, temp_lock_file):
"""Test decorator acquires and releases lock."""
@requires_lock("decorated")
def my_func():
content = temp_lock_file.read_text()
assert "decorated" in content
return "success"
result = my_func()
assert result == "success"
def test_decorator_preserves_metadata(self):
"""Test decorator preserves function metadata."""
@requires_lock("test")
def my_function_with_doc():
"""My docstring."""
assert my_function_with_doc.__name__ == "my_function_with_doc"
assert my_function_with_doc.__doc__ == "My docstring."
def test_decorator_passes_arguments(self, temp_lock_file):
"""Test decorator passes args and kwargs correctly."""
@requires_lock("test")
def func_with_args(a, b, c=None):
return (a, b, c)
result = func_with_args(1, 2, c=3)
assert result == (1, 2, 3)
def test_decorator_handles_lock_error(self, tmp_path):
"""Test decorator handles LockError and exits."""
import typer
test_lock = tmp_path / "dropkit.lock"
ready_file = tmp_path / "ready"
# Script that holds lock
holder_script = f"""
import fcntl
import os
import time
from pathlib import Path
lock_path = "{test_lock}"
ready_path = "{ready_file}"
fd = os.open(lock_path, os.O_RDWR | os.O_CREAT, 0o600)
fcntl.flock(fd, fcntl.LOCK_EX)
Path(ready_path).touch()
time.sleep(30)
"""
p = subprocess.Popen([sys.executable, "-c", holder_script])
try:
# Wait for holder to acquire lock
for _ in range(50):
if ready_file.exists():
break
time.sleep(0.1)
else:
pytest.fail("Holder did not acquire lock in time")
with (
patch("dropkit.lock.LOCK_FILE", test_lock),
patch("dropkit.lock.DEFAULT_TIMEOUT", 0.5),
):
@requires_lock("blocked")
def blocked_func():
return "should not reach"
# Should raise typer.Exit due to LockError
with pytest.raises(typer.Exit) as exc_info:
blocked_func()
assert exc_info.value.exit_code == 1
finally:
p.terminate()
p.wait()
class TestEdgeCases:
"""Tests for edge cases and error conditions."""
def test_default_timeout_value(self):
"""Test default timeout is 30 seconds."""
assert DEFAULT_TIMEOUT == 30.0
def test_lock_file_path(self):
"""Test lock file path is /tmp/dropkit.lock."""
assert Path("/tmp/dropkit.lock") == LOCK_FILE
def test_lock_file_created_with_permissions(self, tmp_path):
"""Test lock file is created with restrictive permissions."""
test_lock = tmp_path / "dropkit.lock"
with patch("dropkit.lock.LOCK_FILE", test_lock):
with operation_lock("test"):
pass
# Check permissions (0o600 = owner read/write only)
mode = test_lock.stat().st_mode & 0o777
assert mode == 0o600
def test_multiple_sequential_locks(self, temp_lock_file):
"""Test acquiring lock multiple times sequentially."""
for i in range(5):
with operation_lock(f"command_{i}"):
content = temp_lock_file.read_text()
assert f"command_{i}" in content
tests/test_main_helpers.py
+360
View File
@@ -0,0 +1,360 @@
"""Tests for main module helper functions."""
from pathlib import Path
from unittest.mock import MagicMock, patch
import pytest
from dropkit.main import (
add_temporary_ssh_rule,
build_droplet_tags,
find_snapshot_action,
get_droplet_name_from_snapshot,
get_snapshot_name,
get_ssh_hostname,
get_user_tag,
is_droplet_tailscale_locked,
prepare_for_hibernate,
)
class TestGetSnapshotName:
"""Tests for get_snapshot_name function."""
def test_simple_name(self):
"""Test snapshot name for simple droplet name."""
assert get_snapshot_name("myvm") == "dropkit-myvm"
def test_name_with_hyphen(self):
"""Test snapshot name for droplet name with hyphen."""
assert get_snapshot_name("my-droplet") == "dropkit-my-droplet"
def test_name_with_numbers(self):
"""Test snapshot name for droplet name with numbers."""
assert get_snapshot_name("test123") == "dropkit-test123"
def test_empty_name(self):
"""Test snapshot name for empty droplet name."""
assert get_snapshot_name("") == "dropkit-"
class TestGetDropletNameFromSnapshot:
"""Tests for get_droplet_name_from_snapshot function."""
def test_valid_dropkit_snapshot(self):
"""Test extracting droplet name from valid dropkit snapshot."""
assert get_droplet_name_from_snapshot("dropkit-myvm") == "myvm"
def test_snapshot_with_hyphen_in_name(self):
"""Test extracting droplet name with hyphens."""
assert get_droplet_name_from_snapshot("dropkit-my-droplet") == "my-droplet"
def test_non_dropkit_snapshot(self):
"""Test with non-dropkit snapshot name."""
assert get_droplet_name_from_snapshot("other-snapshot") is None
def test_partial_prefix(self):
"""Test with partial prefix (should not match)."""
assert get_droplet_name_from_snapshot("dropkit") is None
def test_different_prefix(self):
"""Test with different prefix."""
assert get_droplet_name_from_snapshot("snapshot-myvm") is None
def test_empty_after_prefix(self):
"""Test snapshot name that is just the prefix."""
assert get_droplet_name_from_snapshot("dropkit-") == ""
class TestGetSshHostname:
"""Tests for get_ssh_hostname function."""
def test_simple_name(self):
"""Test SSH hostname for simple droplet name."""
assert get_ssh_hostname("myvm") == "dropkit.myvm"
def test_name_with_hyphen(self):
"""Test SSH hostname for droplet name with hyphen."""
assert get_ssh_hostname("my-droplet") == "dropkit.my-droplet"
class TestGetUserTag:
"""Tests for get_user_tag function."""
def test_simple_username(self):
"""Test user tag for simple username."""
assert get_user_tag("john") == "owner:john"
def test_username_with_underscore(self):
"""Test user tag for username with underscore."""
assert get_user_tag("john_doe") == "owner:john_doe"
class TestBuildDropletTags:
"""Tests for build_droplet_tags function."""
def test_no_extra_tags(self):
"""Test building tags without extra tags."""
tags = build_droplet_tags("john")
assert tags == ["owner:john", "firewall"]
def test_with_extra_tags(self):
"""Test building tags with extra tags."""
tags = build_droplet_tags("john", ["production", "webserver"])
assert tags == ["owner:john", "firewall", "production", "webserver"]
def test_extra_tags_no_duplicates(self):
"""Test that duplicate tags are not added."""
tags = build_droplet_tags("john", ["firewall", "production"])
assert tags == ["owner:john", "firewall", "production"]
def test_empty_extra_tags(self):
"""Test with empty extra tags list."""
tags = build_droplet_tags("john", [])
assert tags == ["owner:john", "firewall"]
def test_none_extra_tags(self):
"""Test with None extra tags."""
tags = build_droplet_tags("john", None)
assert tags == ["owner:john", "firewall"]
class TestFindSnapshotAction:
"""Tests for find_snapshot_action function."""
def test_finds_snapshot_action(self):
"""Test finding a snapshot action in the actions list."""
mock_api = MagicMock()
mock_api.list_droplet_actions.return_value = [
{"id": 1, "type": "power_off", "status": "completed"},
{"id": 2, "type": "snapshot", "status": "in-progress"},
{"id": 3, "type": "power_on", "status": "completed"},
]
result = find_snapshot_action(mock_api, 12345)
assert result is not None
assert result["id"] == 2
assert result["type"] == "snapshot"
mock_api.list_droplet_actions.assert_called_once_with(12345)
def test_returns_first_snapshot_action(self):
"""Test that it returns the first (most recent) snapshot action."""
mock_api = MagicMock()
mock_api.list_droplet_actions.return_value = [
{"id": 10, "type": "snapshot", "status": "in-progress"},
{"id": 5, "type": "snapshot", "status": "completed"},
]
result = find_snapshot_action(mock_api, 12345)
assert result is not None
assert result["id"] == 10
def test_returns_none_when_no_snapshot_action(self):
"""Test returning None when no snapshot action exists."""
mock_api = MagicMock()
mock_api.list_droplet_actions.return_value = [
{"id": 1, "type": "power_off", "status": "completed"},
{"id": 2, "type": "power_on", "status": "completed"},
]
result = find_snapshot_action(mock_api, 12345)
assert result is None
def test_returns_none_when_empty_actions(self):
"""Test returning None when actions list is empty."""
mock_api = MagicMock()
mock_api.list_droplet_actions.return_value = []
result = find_snapshot_action(mock_api, 12345)
assert result is None
@pytest.fixture
def temp_ssh_config(tmp_path):
"""Create a temporary SSH config file."""
ssh_dir = tmp_path / ".ssh"
ssh_dir.mkdir(mode=0o700)
config_path = ssh_dir / "config"
return str(config_path)
class TestIsDropletTailscaleLocked:
"""Tests for is_droplet_tailscale_locked function."""
def test_tailscale_ip_returns_true(self, temp_ssh_config):
"""Test returns True when SSH config has Tailscale IP."""
# Create SSH config with Tailscale IP
Path(temp_ssh_config).write_text("""Host dropkit.myvm
HostName 100.80.123.45
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
result = is_droplet_tailscale_locked(mock_config, "myvm")
assert result is True
def test_public_ip_returns_false(self, temp_ssh_config):
"""Test returns False when SSH config has public IP."""
# Create SSH config with public IP
Path(temp_ssh_config).write_text("""Host dropkit.myvm
HostName 192.168.1.100
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
result = is_droplet_tailscale_locked(mock_config, "myvm")
assert result is False
def test_missing_entry_returns_false(self, temp_ssh_config):
"""Test returns False when SSH config has no entry for droplet."""
# Create SSH config without the target host
Path(temp_ssh_config).write_text("""Host dropkit.othervm
HostName 100.80.123.45
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
result = is_droplet_tailscale_locked(mock_config, "myvm")
assert result is False
def test_nonexistent_config_file_returns_false(self, temp_ssh_config):
"""Test returns False when SSH config file doesn't exist."""
mock_config = MagicMock()
mock_config.ssh.config_path = "/nonexistent/path/config"
result = is_droplet_tailscale_locked(mock_config, "myvm")
assert result is False
class TestAddTemporarySshRule:
"""Tests for add_temporary_ssh_rule function."""
@patch("dropkit.main.subprocess.run")
def test_success(self, mock_run):
"""Test successful SSH rule addition."""
mock_run.return_value = MagicMock(returncode=0, stdout="", stderr="")
result = add_temporary_ssh_rule("dropkit.myvm")
assert result is True
mock_run.assert_called_once()
call_args = mock_run.call_args
assert "dropkit.myvm" in call_args[0][0]
assert "sudo ufw allow in on eth0 to any port 22" in call_args[0][0]
@patch("dropkit.main.subprocess.run")
def test_failure(self, mock_run):
"""Test failed SSH rule addition."""
mock_run.return_value = MagicMock(returncode=1, stdout="", stderr="error")
result = add_temporary_ssh_rule("dropkit.myvm")
assert result is False
@patch("dropkit.main.subprocess.run")
def test_timeout(self, mock_run):
"""Test SSH timeout."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 30)
result = add_temporary_ssh_rule("dropkit.myvm")
assert result is False
class TestPrepareForHibernate:
"""Tests for prepare_for_hibernate function."""
def test_not_tailscale_locked_returns_false(self, temp_ssh_config):
"""Test returns False when droplet is not Tailscale locked."""
# Create SSH config with public IP (not Tailscale locked)
Path(temp_ssh_config).write_text("""Host dropkit.myvm
HostName 192.168.1.100
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
mock_api = MagicMock()
mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "192.168.1.100"}]}}
result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm")
assert result is False
@patch("dropkit.main.tailscale_logout")
@patch("dropkit.main.add_temporary_ssh_rule")
@patch("dropkit.main.add_ssh_host")
def test_tailscale_locked_returns_true(
self, mock_add_ssh_host, mock_add_temp_rule, mock_logout, temp_ssh_config
):
"""Test returns True when droplet is Tailscale locked."""
# Create SSH config with Tailscale IP
Path(temp_ssh_config).write_text("""Host dropkit.myvm
HostName 100.80.123.45
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
mock_config.ssh.identity_file = "~/.ssh/id_ed25519"
mock_api = MagicMock()
mock_api.get_username.return_value = "testuser"
mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "203.0.113.50"}]}}
mock_add_temp_rule.return_value = True
mock_logout.return_value = True
result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm")
assert result is True
mock_add_temp_rule.assert_called_once()
mock_logout.assert_called_once()
mock_add_ssh_host.assert_called_once()
@patch("dropkit.main.tailscale_logout")
@patch("dropkit.main.add_temporary_ssh_rule")
@patch("dropkit.main.add_ssh_host")
def test_temp_rule_failure_skips_logout(
self, mock_add_ssh_host, mock_add_temp_rule, mock_logout, temp_ssh_config
):
"""Test returns True but skips logout if temp rule fails (safety)."""
# Create SSH config with Tailscale IP
Path(temp_ssh_config).write_text("""Host dropkit.myvm
HostName 100.80.123.45
User ubuntu
""")
mock_config = MagicMock()
mock_config.ssh.config_path = temp_ssh_config
mock_config.ssh.identity_file = "~/.ssh/id_ed25519"
mock_api = MagicMock()
mock_api.get_username.return_value = "testuser"
mock_droplet = {"networks": {"v4": [{"type": "public", "ip_address": "203.0.113.50"}]}}
mock_add_temp_rule.return_value = False # Simulating failure
result = prepare_for_hibernate(mock_config, mock_api, mock_droplet, "myvm")
# Should still return True because we detected Tailscale lockdown
assert result is True
# But logout should NOT be called (safety - need public IP fallback first)
mock_logout.assert_not_called()
# And SSH config should NOT be updated (early return)
mock_add_ssh_host.assert_not_called()
tests/test_rename.py
+354
View File
@@ -0,0 +1,354 @@
"""Tests for the rename command functionality."""
from unittest.mock import MagicMock, patch
from typer.testing import CliRunner
from dropkit.main import app
runner = CliRunner()
def create_mock_droplet(
droplet_id: int = 12345,
name: str = "test-droplet",
status: str = "active",
ip_address: str = "192.168.1.100",
) -> dict:
"""Create a mock droplet dictionary."""
return {
"id": droplet_id,
"name": name,
"status": status,
"networks": {
"v4": [
{"type": "public", "ip_address": ip_address},
{"type": "private", "ip_address": "10.0.0.1"},
]
},
}
def create_mock_config():
"""Create a mock config object."""
mock_config = MagicMock()
mock_config.ssh.config_path = "~/.ssh/config"
mock_config.ssh.identity_file = "~/.ssh/id_ed25519"
return mock_config
class TestRenameCommand:
"""Tests for the rename command."""
@patch("dropkit.main.add_ssh_host")
@patch("dropkit.main.remove_ssh_host")
@patch("dropkit.main.host_exists")
@patch("dropkit.main.Prompt.ask")
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_successful_rename(
self,
mock_load_config,
mock_find_droplet,
mock_prompt,
mock_host_exists,
mock_remove_ssh,
mock_add_ssh,
):
"""Test successful droplet rename."""
# Setup mocks
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="old-name", status="active")
mock_find_droplet.return_value = (mock_droplet, "testuser")
# No existing droplets with new name
mock_api.list_droplets.return_value = [mock_droplet]
mock_api.get_droplet.return_value = mock_droplet
mock_api.rename_droplet.return_value = {"id": 99999}
mock_api.wait_for_action_complete.return_value = {"status": "completed"}
mock_prompt.return_value = "yes"
mock_host_exists.return_value = True
# Run command
result = runner.invoke(app, ["rename", "old-name", "new-name"])
# Verify
assert result.exit_code == 0
assert "successfully renamed" in result.output.lower()
mock_api.rename_droplet.assert_called_once_with(12345, "new-name")
mock_remove_ssh.assert_called_once()
mock_add_ssh.assert_called_once()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_same_name_exits_early(self, mock_load_config, mock_find_droplet):
"""Test that renaming to the same name exits early."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="test-droplet")
mock_find_droplet.return_value = (mock_droplet, "testuser")
result = runner.invoke(app, ["rename", "test-droplet", "test-droplet"])
assert result.exit_code == 0
assert "same as current name" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_droplet_not_found(self, mock_load_config, mock_find_droplet):
"""Test rename when droplet is not found."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_find_droplet.return_value = (None, "testuser")
result = runner.invoke(app, ["rename", "nonexistent", "new-name"])
assert result.exit_code == 1
assert "not found" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_name_conflict(self, mock_load_config, mock_find_droplet):
"""Test rename fails when new name already exists for the same user."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(droplet_id=12345, name="old-name")
mock_find_droplet.return_value = (mock_droplet, "testuser")
# Another droplet with the target name already exists
existing_droplet = create_mock_droplet(droplet_id=99999, name="existing-name")
mock_api.list_droplets.return_value = [mock_droplet, existing_droplet]
result = runner.invoke(app, ["rename", "old-name", "existing-name"])
assert result.exit_code == 1
assert "already exists" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_droplet_not_active(self, mock_load_config, mock_find_droplet):
"""Test rename fails when droplet is not active (powered off)."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
# Droplet is off
mock_droplet = create_mock_droplet(name="old-name", status="off")
mock_find_droplet.return_value = (mock_droplet, "testuser")
mock_api.list_droplets.return_value = [mock_droplet]
result = runner.invoke(app, ["rename", "old-name", "new-name"])
assert result.exit_code == 1
assert "off" in result.output.lower()
assert "dropkit on" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_droplet_status_new(self, mock_load_config, mock_find_droplet):
"""Test rename fails when droplet status is 'new' (still initializing)."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="old-name", status="new")
mock_find_droplet.return_value = (mock_droplet, "testuser")
mock_api.list_droplets.return_value = [mock_droplet]
result = runner.invoke(app, ["rename", "old-name", "new-name"])
assert result.exit_code == 1
assert "new" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.Prompt.ask")
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_user_cancels(self, mock_load_config, mock_find_droplet, mock_prompt):
"""Test rename cancelled when user says no."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="old-name", status="active")
mock_find_droplet.return_value = (mock_droplet, "testuser")
mock_api.list_droplets.return_value = [mock_droplet]
mock_api.get_droplet.return_value = mock_droplet
mock_prompt.return_value = "no"
result = runner.invoke(app, ["rename", "old-name", "new-name"])
assert result.exit_code == 0
assert "cancelled" in result.output.lower()
mock_api.rename_droplet.assert_not_called()
@patch("dropkit.main.host_exists")
@patch("dropkit.main.Prompt.ask")
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_no_ssh_config_entry(
self, mock_load_config, mock_find_droplet, mock_prompt, mock_host_exists
):
"""Test rename when no SSH config entry exists for old name."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="old-name", status="active")
mock_find_droplet.return_value = (mock_droplet, "testuser")
mock_api.list_droplets.return_value = [mock_droplet]
mock_api.get_droplet.return_value = mock_droplet
mock_api.rename_droplet.return_value = {"id": 99999}
mock_api.wait_for_action_complete.return_value = {"status": "completed"}
mock_prompt.return_value = "yes"
mock_host_exists.return_value = False # No SSH entry
result = runner.invoke(app, ["rename", "old-name", "new-name"])
assert result.exit_code == 0
assert "successfully renamed" in result.output.lower()
assert "not found" in result.output.lower() # SSH config not found message
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_only_checks_user_droplets(self, mock_load_config, mock_find_droplet):
"""Test that name conflict check only considers user's own droplets."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(droplet_id=12345, name="old-name")
mock_find_droplet.return_value = (mock_droplet, "testuser")
# Only user's droplet is returned (filtered by tag)
mock_api.list_droplets.return_value = [mock_droplet]
# Existing droplet with same name but different user won't be in list
# because list_droplets is called with tag_name=owner:testuser
runner.invoke(app, ["rename", "old-name", "target-name"])
# Should proceed past name conflict check
# Verify list_droplets was called with the user tag
mock_api.list_droplets.assert_called_once()
call_kwargs = mock_api.list_droplets.call_args
assert "owner:testuser" in str(call_kwargs)
class TestRenameEdgeCases:
"""Edge case tests for rename command."""
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_droplet_missing_id(self, mock_load_config, mock_find_droplet):
"""Test rename when droplet has no ID."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
# Droplet without ID
mock_droplet = {"name": "old-name", "status": "active"}
mock_find_droplet.return_value = (mock_droplet, "testuser")
result = runner.invoke(app, ["rename", "old-name", "new-name"])
assert result.exit_code == 1
assert "could not determine droplet id" in result.output.lower()
@patch("dropkit.main.add_ssh_host")
@patch("dropkit.main.remove_ssh_host")
@patch("dropkit.main.host_exists")
@patch("dropkit.main.Prompt.ask")
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_ssh_update_failure_continues(
self,
mock_load_config,
mock_find_droplet,
mock_prompt,
mock_host_exists,
mock_remove_ssh,
mock_add_ssh,
):
"""Test that SSH config update failure doesn't block rename success."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
mock_droplet = create_mock_droplet(name="old-name", status="active")
mock_find_droplet.return_value = (mock_droplet, "testuser")
mock_api.list_droplets.return_value = [mock_droplet]
mock_api.get_droplet.return_value = mock_droplet
mock_api.rename_droplet.return_value = {"id": 99999}
mock_api.wait_for_action_complete.return_value = {"status": "completed"}
mock_prompt.return_value = "yes"
mock_host_exists.return_value = True
mock_remove_ssh.side_effect = Exception("SSH config error")
result = runner.invoke(app, ["rename", "old-name", "new-name"])
# Rename should still succeed even if SSH config update fails
assert result.exit_code == 0
assert "successfully renamed" in result.output.lower()
assert "could not update ssh config" in result.output.lower()
@patch("dropkit.main.find_user_droplet")
@patch("dropkit.main.load_config_and_api")
def test_rename_multiple_droplets_same_user(self, mock_load_config, mock_find_droplet):
"""Test name conflict when user has multiple droplets."""
mock_api = MagicMock()
mock_config_manager = MagicMock()
mock_config_manager.config = create_mock_config()
mock_load_config.return_value = (mock_config_manager, mock_api)
droplet_to_rename = create_mock_droplet(droplet_id=111, name="droplet-a")
other_droplet = create_mock_droplet(droplet_id=222, name="droplet-b")
target_droplet = create_mock_droplet(droplet_id=333, name="droplet-c")
mock_find_droplet.return_value = (droplet_to_rename, "testuser")
# User has 3 droplets
mock_api.list_droplets.return_value = [
droplet_to_rename,
other_droplet,
target_droplet,
]
# Try to rename to existing name
result = runner.invoke(app, ["rename", "droplet-a", "droplet-c"])
assert result.exit_code == 1
assert "already exists" in result.output.lower()
tests/test_ssh_config.py
+985
View File
@@ -0,0 +1,985 @@
"""Comprehensive tests for SSH config management."""
from pathlib import Path
import pytest
from dropkit.ssh_config import (
add_ssh_host,
get_ssh_host_ip,
remove_known_hosts_entry,
remove_ssh_host,
)
@pytest.fixture
def temp_ssh_dir(tmp_path):
"""Create a temporary SSH directory for testing."""
ssh_dir = tmp_path / ".ssh"
ssh_dir.mkdir(mode=0o700)
yield ssh_dir
# Cleanup is automatic with tmp_path
@pytest.fixture
def temp_config(temp_ssh_dir):
"""Create a temporary SSH config file path."""
config_path = temp_ssh_dir / "config"
return str(config_path)
class TestAddSSHHost:
"""Tests for add_ssh_host function."""
def test_empty_file(self, temp_config):
"""Test adding host to empty SSH config file."""
# Create empty file
Path(temp_config).touch(mode=0o600)
add_ssh_host(temp_config, "myhost", "192.168.1.1", "ubuntu")
content = Path(temp_config).read_text()
assert "Host myhost" in content
assert "HostName 192.168.1.1" in content
assert "User ubuntu" in content
def test_non_existent_file(self, temp_config):
"""Test creating new file and directory if they don't exist."""
# Ensure file doesn't exist
assert not Path(temp_config).exists()
add_ssh_host(temp_config, "myhost", "192.168.1.1", "ubuntu", "~/.ssh/id_ed25519")
assert Path(temp_config).exists()
content = Path(temp_config).read_text()
assert "Host myhost" in content
assert "IdentityFile ~/.ssh/id_ed25519" in content
def test_simple_addition(self, temp_config):
"""Test adding host to config with existing unrelated hosts."""
# Create config with existing host
existing = """Host other-server
HostName 10.0.0.1
User admin
Host another-host
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "myhost", "192.168.1.1", "ubuntu")
content = Path(temp_config).read_text()
assert "Host other-server" in content
assert "Host another-host" in content
assert "Host myhost" in content
# Count "Host " with space to avoid matching "HostName"
assert content.count("Host ") == 3
def test_update_existing_host(self, temp_config):
"""Test replacing existing host entry without duplication."""
# Create config with existing host
existing = """Host myhost
HostName 10.0.0.1
User old_user
Host other-host
HostName 10.0.0.2
User admin
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "myhost", "192.168.1.1", "new_user")
content = Path(temp_config).read_text()
assert content.count("Host myhost") == 1 # Should not duplicate
assert "HostName 192.168.1.1" in content
assert "User new_user" in content
assert "old_user" not in content
assert "10.0.0.1" not in content
assert "Host other-host" in content # Other hosts preserved
def test_host_with_similar_names(self, temp_config):
"""Test that similar host names don't interfere."""
existing = """Host myhost
HostName 10.0.0.1
User user1
Host myhost-dev
HostName 10.0.0.2
User user2
Host myhost-prod
HostName 10.0.0.3
User user3
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "myhost", "192.168.1.1", "newuser")
content = Path(temp_config).read_text()
assert "Host myhost-dev" in content
assert "10.0.0.2" in content
assert "Host myhost-prod" in content
assert "10.0.0.3" in content
# myhost should be updated
lines = content.split("\n")
myhost_line = [i for i, line in enumerate(lines) if line.strip() == "Host myhost"][-1]
# Check that the next few lines after myhost contain the new values
section = "\n".join(lines[myhost_line : myhost_line + 5])
assert "192.168.1.1" in section
assert "newuser" in section
def test_host_as_substring(self, temp_config):
"""Test that substring hosts are handled correctly."""
existing = """Host prod-server
HostName 10.0.0.1
User admin
Host prod
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "prod", "192.168.1.1", "ubuntu")
content = Path(temp_config).read_text()
assert "Host prod-server" in content
assert "10.0.0.1" in content
# prod should be updated
assert content.count("Host prod\n") >= 1
def test_multiple_blank_lines(self, temp_config):
"""Test handling config with various blank lines."""
existing = """Host server1
HostName 10.0.0.1
User admin
Host server2
HostName 10.0.0.2
User root
Host server3
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "user")
content = Path(temp_config).read_text()
assert "Host server1" in content
assert "Host server2" in content
assert "Host server3" in content
assert "Host newhost" in content
def test_comments_in_config(self, temp_config):
"""Test preserving comments in config."""
existing = """# This is a comment
Host server1
HostName 10.0.0.1
User admin
# Another comment
# Global settings
Host server2
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "user")
content = Path(temp_config).read_text()
assert "# This is a comment" in content
assert "# Another comment" in content
assert "# Global settings" in content
def test_mixed_indentation(self, temp_config):
"""Test handling tabs and spaces."""
existing = """Host server1
HostName 10.0.0.1
User admin
Host server2
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "user")
content = Path(temp_config).read_text()
# Original hosts should be preserved
assert "Host server1" in content
assert "Host server2" in content
# New host should be added
assert "Host newhost" in content
def test_host_at_beginning(self, temp_config):
"""Test updating a host that's the first entry."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host second-host
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "first-host", "192.168.1.1", "newuser")
content = Path(temp_config).read_text()
assert content.count("Host first-host") == 1
assert "192.168.1.1" in content
assert "10.0.0.1" not in content
assert "Host second-host" in content
def test_host_at_end(self, temp_config):
"""Test updating a host that's the last entry."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host last-host
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "last-host", "192.168.1.1", "newuser")
content = Path(temp_config).read_text()
assert "Host first-host" in content
assert content.count("Host last-host") == 1
assert "192.168.1.1" in content
assert "10.0.0.2" not in content
def test_host_in_middle(self, temp_config):
"""Test updating a host between others."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host middle-host
HostName 10.0.0.2
User root
Host last-host
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "middle-host", "192.168.1.1", "newuser")
content = Path(temp_config).read_text()
assert "Host first-host" in content
assert "Host last-host" in content
assert content.count("Host middle-host") == 1
assert "192.168.1.1" in content
assert "10.0.0.2" not in content
def test_complex_host_entries(self, temp_config):
"""Test host with many configuration options."""
existing = """Host complex-host
HostName 10.0.0.1
User admin
Port 2222
ForwardAgent yes
ProxyJump bastion
IdentityFile ~/.ssh/custom_key
StrictHostKeyChecking no
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "complex-host", "192.168.1.1", "newuser", "~/.ssh/new_key")
content = Path(temp_config).read_text()
assert content.count("Host complex-host") == 1
assert "192.168.1.1" in content
assert "newuser" in content
# Old complex options should be replaced
assert "Port 2222" not in content
def test_host_with_extra_spaces(self, temp_config):
"""Test host declaration with multiple spaces."""
existing = """Host server-with-spaces
HostName 10.0.0.1
User admin
"""
Path(temp_config).write_text(existing)
# This won't match because our check is exact
add_ssh_host(temp_config, "server-with-spaces", "192.168.1.1", "newuser")
content = Path(temp_config).read_text()
# Should add a new entry since "Host server-with-spaces" != "Host server-with-spaces"
assert "Host server-with-spaces" in content
def test_host_patterns(self, temp_config):
"""Test existing wildcard hosts."""
existing = """Host *.example.com
User admin
IdentityFile ~/.ssh/example_key
Host myserver
HostName 10.0.0.1
User root
"""
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "ubuntu")
content = Path(temp_config).read_text()
assert "Host *.example.com" in content
assert "Host myserver" in content
assert "Host newhost" in content
def test_no_trailing_newline(self, temp_config):
"""Test config file without final newline."""
existing = """Host myhost
HostName 10.0.0.1
User admin""" # No trailing newline
Path(temp_config).write_text(existing)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "ubuntu")
content = Path(temp_config).read_text()
assert "Host myhost" in content
assert "Host newhost" in content
# Verify Host newhost is on its own line (not concatenated to previous line)
lines = content.split("\n")
host_lines = [line for line in lines if "Host newhost" in line]
assert len(host_lines) == 1
assert host_lines[0].strip() == "Host newhost"
def test_permissions_preserved(self, temp_config):
"""Test that file permissions are set correctly."""
Path(temp_config).touch(mode=0o644)
add_ssh_host(temp_config, "myhost", "192.168.1.1", "ubuntu")
# File should now have 0600 permissions
mode = Path(temp_config).stat().st_mode & 0o777
assert mode == 0o600
def test_backup_created(self, temp_config):
"""Test that backup file is created."""
existing = """Host oldhost
HostName 10.0.0.1
User admin
"""
Path(temp_config).write_text(existing)
Path(temp_config).chmod(0o600)
add_ssh_host(temp_config, "newhost", "192.168.1.1", "ubuntu")
backup_path = Path(temp_config).parent / "config.bak"
assert backup_path.exists()
backup_content = backup_path.read_text()
assert "Host oldhost" in backup_content
assert "Host newhost" not in backup_content
# Check backup has same permissions
backup_mode = backup_path.stat().st_mode & 0o777
assert backup_mode == 0o600
class TestRemoveSSHHost:
"""Tests for remove_ssh_host function."""
def test_non_existent_file(self, temp_config):
"""Test removing from non-existent config."""
assert not Path(temp_config).exists()
result = remove_ssh_host(temp_config, "myhost")
assert result is False
def test_host_doesnt_exist(self, temp_config):
"""Test removing host that's not in config."""
existing = """Host server1
HostName 10.0.0.1
User admin
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "nonexistent")
assert result is False
content = Path(temp_config).read_text()
assert "Host server1" in content # Should be unchanged
def test_remove_first_host(self, temp_config):
"""Test removing the first host entry."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host second-host
HostName 10.0.0.2
User root
Host third-host
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "first-host")
assert result is True
content = Path(temp_config).read_text()
assert "Host first-host" not in content
assert "Host second-host" in content
assert "Host third-host" in content
def test_remove_last_host(self, temp_config):
"""Test removing the last host entry."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host second-host
HostName 10.0.0.2
User root
Host last-host
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "last-host")
assert result is True
content = Path(temp_config).read_text()
assert "Host first-host" in content
assert "Host second-host" in content
assert "Host last-host" not in content
def test_remove_middle_host(self, temp_config):
"""Test removing a host between others."""
existing = """Host first-host
HostName 10.0.0.1
User admin
Host middle-host
HostName 10.0.0.2
User root
Host last-host
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "middle-host")
assert result is True
content = Path(temp_config).read_text()
assert "Host first-host" in content
assert "Host middle-host" not in content
assert "Host last-host" in content
def test_remove_only_host(self, temp_config):
"""Test removing the only host in config."""
existing = """Host only-host
HostName 10.0.0.1
User admin
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "only-host")
assert result is True
content = Path(temp_config).read_text()
assert "Host only-host" not in content
# File should be essentially empty or just whitespace
assert content.strip() == "" or "Host" not in content
def test_remove_similar_host_names(self, temp_config):
"""Test removing specific host when similar names exist."""
existing = """Host myhost
HostName 10.0.0.1
User user1
Host myhost-dev
HostName 10.0.0.2
User user2
Host myhost-prod
HostName 10.0.0.3
User user3
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "myhost")
assert result is True
content = Path(temp_config).read_text()
assert "Host myhost\n" not in content
assert "10.0.0.1" not in content
assert "Host myhost-dev" in content
assert "10.0.0.2" in content
assert "Host myhost-prod" in content
assert "10.0.0.3" in content
def test_remove_host_with_many_options(self, temp_config):
"""Test removing host with many configuration lines."""
existing = """Host simple-host
HostName 10.0.0.1
User admin
Host complex-host
HostName 10.0.0.2
User root
Port 2222
ForwardAgent yes
ProxyJump bastion
IdentityFile ~/.ssh/custom_key
StrictHostKeyChecking no
LocalForward 8080 localhost:80
Host another-host
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "complex-host")
assert result is True
content = Path(temp_config).read_text()
assert "Host simple-host" in content
assert "Host complex-host" not in content
assert "Port 2222" not in content
assert "ProxyJump bastion" not in content
assert "Host another-host" in content
def test_backup_created_on_remove(self, temp_config):
"""Test that backup is created when removing host."""
existing = """Host myhost
HostName 10.0.0.1
User admin
Host otherhost
HostName 10.0.0.2
User root
"""
Path(temp_config).write_text(existing)
Path(temp_config).chmod(0o600)
remove_ssh_host(temp_config, "myhost")
backup_path = Path(temp_config).parent / "config.bak"
assert backup_path.exists()
backup_content = backup_path.read_text()
assert "Host myhost" in backup_content
assert "Host otherhost" in backup_content
# Check backup has same permissions
backup_mode = backup_path.stat().st_mode & 0o777
assert backup_mode == 0o600
def test_remove_preserves_comments(self, temp_config):
"""Test that comments are preserved when removing host."""
existing = """# Global comment
Host keephost
HostName 10.0.0.1
User admin
# This host will be removed
Host removehost
HostName 10.0.0.2
User root
# Another comment
Host anotherhost
HostName 10.0.0.3
User ubuntu
"""
Path(temp_config).write_text(existing)
result = remove_ssh_host(temp_config, "removehost")
assert result is True
content = Path(temp_config).read_text()
assert "# Global comment" in content
assert "# Another comment" in content
assert "Host keephost" in content
assert "Host removehost" not in content
assert "Host anotherhost" in content
class TestGetSSHHostIP:
"""Tests for get_ssh_host_ip function."""
def test_valid_host(self, temp_config):
"""Test getting IP for existing host."""
existing = """Host myhost
HostName 192.168.1.100
User ubuntu
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "myhost")
assert result == "192.168.1.100"
def test_host_not_found(self, temp_config):
"""Test getting IP for non-existent host."""
existing = """Host otherhost
HostName 192.168.1.100
User ubuntu
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "myhost")
assert result is None
def test_missing_hostname_field(self, temp_config):
"""Test host entry without HostName field."""
existing = """Host myhost
User ubuntu
IdentityFile ~/.ssh/id_rsa
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "myhost")
assert result is None
def test_non_existent_file(self, temp_config):
"""Test with non-existent config file."""
result = get_ssh_host_ip(temp_config, "myhost")
assert result is None
def test_tailscale_ip(self, temp_config):
"""Test getting Tailscale IP address."""
existing = """Host dropkit.myhost
HostName 100.80.123.45
User ubuntu
ForwardAgent yes
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "dropkit.myhost")
assert result == "100.80.123.45"
def test_multiple_hosts(self, temp_config):
"""Test getting IP from config with multiple hosts."""
existing = """Host firsthost
HostName 10.0.0.1
User admin
Host targethost
HostName 192.168.1.50
User ubuntu
Host thirdhost
HostName 10.0.0.3
User root
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "targethost")
assert result == "192.168.1.50"
def test_host_with_extra_whitespace(self, temp_config):
"""Test HostName with extra whitespace."""
existing = """Host myhost
HostName 192.168.1.100
User ubuntu
"""
Path(temp_config).write_text(existing)
result = get_ssh_host_ip(temp_config, "myhost")
assert result == "192.168.1.100"
def test_multiple_hosts_on_same_line(self, temp_config):
"""Test host with multiple aliases on same line."""
existing = """Host myhost myalias anotherhost
HostName 192.168.1.100
User ubuntu
"""
Path(temp_config).write_text(existing)
# Should work for any of the aliases
assert get_ssh_host_ip(temp_config, "myhost") == "192.168.1.100"
assert get_ssh_host_ip(temp_config, "myalias") == "192.168.1.100"
assert get_ssh_host_ip(temp_config, "anotherhost") == "192.168.1.100"
def test_hostname_vs_host(self, temp_config):
"""Test that we distinguish between Host directive and HostName option."""
existing = """Host realhost
HostName 192.168.1.100
User ubuntu
Host anotherreal
HostName 10.0.0.1
User admin
"""
Path(temp_config).write_text(existing)
# Should not find HostName as a host alias
assert get_ssh_host_ip(temp_config, "192.168.1.100") is None
# Should find the actual hosts
assert get_ssh_host_ip(temp_config, "realhost") == "192.168.1.100"
class TestRemoveKnownHostsEntry:
"""Tests for remove_known_hosts_entry function."""
@pytest.fixture
def temp_known_hosts(self, temp_ssh_dir):
"""Create a temporary known_hosts file path."""
known_hosts_path = temp_ssh_dir / "known_hosts"
return str(known_hosts_path)
def test_non_existent_file(self, temp_known_hosts):
"""Test removing from non-existent known_hosts."""
assert not Path(temp_known_hosts).exists()
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 0
def test_no_matching_entries(self, temp_known_hosts):
"""Test removing hostname that's not in known_hosts."""
existing = """otherhost ssh-ed25519 AAAA...
192.168.1.1 ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 0
content = Path(temp_known_hosts).read_text()
assert "otherhost" in content
assert "192.168.1.1" in content
def test_remove_single_hostname(self, temp_known_hosts):
"""Test removing a single hostname entry."""
existing = """myhost ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "myhost" not in content
assert "otherhost" in content
def test_remove_ip_address(self, temp_known_hosts):
"""Test removing an IP address entry."""
existing = """192.168.1.100 ssh-ed25519 AAAA...
10.0.0.1 ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["192.168.1.100"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "192.168.1.100" not in content
assert "10.0.0.1" in content
def test_remove_multiple_entries(self, temp_known_hosts):
"""Test removing hostname and IP address together."""
existing = """tobcloud.myhost ssh-ed25519 AAAA...
100.80.123.45 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["tobcloud.myhost", "100.80.123.45"])
assert result == 2
content = Path(temp_known_hosts).read_text()
assert "tobcloud.myhost" not in content
assert "100.80.123.45" not in content
assert "otherhost" in content
def test_comma_separated_hostnames(self, temp_known_hosts):
"""Test removing entry with comma-separated hostnames."""
existing = """myhost,192.168.1.100 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
# Should match by hostname
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "myhost" not in content
assert "192.168.1.100" not in content # Whole line removed
assert "otherhost" in content
def test_comma_separated_match_by_ip(self, temp_known_hosts):
"""Test removing entry by IP when comma-separated."""
existing = """myhost,192.168.1.100 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
# Should match by IP
result = remove_known_hosts_entry(temp_known_hosts, ["192.168.1.100"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "myhost" not in content
assert "192.168.1.100" not in content
assert "otherhost" in content
def test_bracketed_entry(self, temp_known_hosts):
"""Test removing bracketed entry like [hostname]:port."""
existing = """[myhost]:2222 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "[myhost]:2222" not in content
assert "otherhost" in content
def test_bracketed_ip_entry(self, temp_known_hosts):
"""Test removing bracketed IP entry like [192.168.1.1]:2222."""
existing = """[192.168.1.100]:2222 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["192.168.1.100"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "192.168.1.100" not in content
assert "otherhost" in content
def test_hashed_entries_preserved(self, temp_known_hosts):
"""Test that hashed entries (|1|...) are preserved."""
existing = """|1|abc123...= ssh-ed25519 AAAA...
myhost ssh-rsa AAAA...
|1|def456...= ssh-ed25519 AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "|1|abc123" in content
assert "|1|def456" in content
assert "myhost" not in content
def test_comments_preserved(self, temp_known_hosts):
"""Test that comments are preserved."""
existing = """# This is a comment
myhost ssh-ed25519 AAAA...
# Another comment
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "# This is a comment" in content
assert "# Another comment" in content
assert "myhost" not in content
assert "otherhost" in content
def test_empty_lines_preserved(self, temp_known_hosts):
"""Test that empty lines are preserved."""
existing = """myhost ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "myhost" not in content
assert "otherhost" in content
# Should preserve structure
assert "\n\n" in content or content.count("\n") >= 2
def test_case_insensitive_matching(self, temp_known_hosts):
"""Test that hostname matching is case-insensitive."""
existing = """MyHost ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost"])
assert result == 1
content = Path(temp_known_hosts).read_text()
assert "MyHost" not in content
assert "otherhost" in content
def test_backup_created(self, temp_known_hosts):
"""Test that backup file is created."""
existing = """myhost ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
Path(temp_known_hosts).chmod(0o600)
remove_known_hosts_entry(temp_known_hosts, ["myhost"])
backup_path = Path(temp_known_hosts).parent / "known_hosts.bak"
assert backup_path.exists()
backup_content = backup_path.read_text()
assert "myhost" in backup_content
assert "otherhost" in backup_content
# Check backup has same permissions
backup_mode = backup_path.stat().st_mode & 0o777
assert backup_mode == 0o600
def test_remove_all_entries(self, temp_known_hosts):
"""Test removing all entries from known_hosts."""
existing = """myhost ssh-ed25519 AAAA...
192.168.1.100 ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(temp_known_hosts, ["myhost", "192.168.1.100"])
assert result == 2
content = Path(temp_known_hosts).read_text()
assert content.strip() == ""
def test_tobcloud_hostname_format(self, temp_known_hosts):
"""Test removing tobcloud-style hostname."""
existing = """tobcloud.my-droplet ssh-ed25519 AAAA...
100.80.123.45 ssh-ed25519 AAAA...
otherhost ssh-rsa AAAA...
"""
Path(temp_known_hosts).write_text(existing)
result = remove_known_hosts_entry(
temp_known_hosts, ["tobcloud.my-droplet", "100.80.123.45"]
)
assert result == 2
content = Path(temp_known_hosts).read_text()
assert "tobcloud.my-droplet" not in content
assert "100.80.123.45" not in content
assert "otherhost" in content
tests/test_tailscale.py
+536
View File
@@ -0,0 +1,536 @@
"""Tests for Tailscale integration functions."""
import json
from unittest.mock import MagicMock, patch
import pytest
from dropkit.config import TailscaleConfig
from dropkit.main import (
check_local_tailscale,
check_tailscale_installed,
install_tailscale_on_droplet,
is_tailscale_ip,
lock_down_to_tailscale,
run_tailscale_up,
setup_tailscale,
tailscale_logout,
verify_tailscale_ssh,
)
class TestIsTailscaleIP:
"""Tests for is_tailscale_ip function."""
def test_valid_tailscale_ip_lower_bound(self):
"""Test lower bound of CGNAT range (100.64.0.0)."""
assert is_tailscale_ip("100.64.0.0") is True
assert is_tailscale_ip("100.64.0.1") is True
def test_valid_tailscale_ip_upper_bound(self):
"""Test upper bound of CGNAT range (100.127.255.255)."""
assert is_tailscale_ip("100.127.255.255") is True
assert is_tailscale_ip("100.127.0.1") is True
def test_valid_tailscale_ip_middle_range(self):
"""Test middle of CGNAT range."""
assert is_tailscale_ip("100.100.50.25") is True
assert is_tailscale_ip("100.80.1.1") is True
def test_invalid_ip_below_cgnat_range(self):
"""Test IPs below CGNAT range (100.0.0.0 - 100.63.255.255)."""
assert is_tailscale_ip("100.0.0.1") is False
assert is_tailscale_ip("100.63.255.255") is False
def test_invalid_ip_above_cgnat_range(self):
"""Test IPs above CGNAT range (100.128.0.0+)."""
assert is_tailscale_ip("100.128.0.0") is False
assert is_tailscale_ip("100.200.1.1") is False
def test_invalid_ip_wrong_first_octet(self):
"""Test IPs with wrong first octet."""
assert is_tailscale_ip("192.168.1.1") is False
assert is_tailscale_ip("10.0.0.1") is False
assert is_tailscale_ip("172.16.0.1") is False
def test_invalid_ip_format_too_few_octets(self):
"""Test invalid IP format with too few octets."""
assert is_tailscale_ip("100.64.0") is False
assert is_tailscale_ip("100.64") is False
assert is_tailscale_ip("100") is False
def test_invalid_ip_format_too_many_octets(self):
"""Test invalid IP format with too many octets."""
assert is_tailscale_ip("100.64.0.1.5") is False
def test_invalid_ip_format_non_numeric(self):
"""Test invalid IP format with non-numeric values."""
assert is_tailscale_ip("100.64.abc.1") is False
assert is_tailscale_ip("foo.bar.baz.qux") is False
def test_invalid_ip_format_out_of_range_octets(self):
"""Test IPs with octets out of 0-255 range."""
assert is_tailscale_ip("100.64.256.1") is False
assert is_tailscale_ip("100.64.0.-1") is False
def test_invalid_ip_empty_string(self):
"""Test empty string."""
assert is_tailscale_ip("") is False
def test_invalid_ip_none_type(self):
"""Test None type (should return False, not raise)."""
assert is_tailscale_ip(None) is False # type: ignore
class TestCheckLocalTailscale:
"""Tests for check_local_tailscale function."""
@patch("dropkit.main.subprocess.run")
def test_tailscale_running(self, mock_run):
"""Test when Tailscale is running locally."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=json.dumps({"BackendState": "Running"}).encode("utf-8"),
)
assert check_local_tailscale() is True
@patch("dropkit.main.subprocess.run")
def test_tailscale_not_running(self, mock_run):
"""Test when Tailscale is installed but not running."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=json.dumps({"BackendState": "Stopped"}).encode("utf-8"),
)
assert check_local_tailscale() is False
@patch("dropkit.main.subprocess.run")
def test_tailscale_command_fails(self, mock_run):
"""Test when tailscale command returns non-zero."""
mock_run.return_value = MagicMock(returncode=1)
assert check_local_tailscale() is False
@patch("dropkit.main.subprocess.run")
def test_tailscale_not_installed(self, mock_run):
"""Test when tailscale is not installed."""
mock_run.side_effect = FileNotFoundError()
assert check_local_tailscale() is False
@patch("dropkit.main.subprocess.run")
def test_tailscale_timeout(self, mock_run):
"""Test when tailscale command times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("tailscale", 5)
assert check_local_tailscale() is False
@patch("dropkit.main.subprocess.run")
def test_invalid_json_response(self, mock_run):
"""Test when tailscale returns invalid JSON."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"not valid json",
)
assert check_local_tailscale() is False
class TestRunTailscaleUp:
"""Tests for run_tailscale_up function."""
@patch("dropkit.main.subprocess.run")
def test_extracts_auth_url(self, mock_run):
"""Test that auth URL is extracted from output."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"To authenticate, visit:\n\n\thttps://login.tailscale.com/a/abc123\n",
)
url = run_tailscale_up("dropkit.test")
assert url == "https://login.tailscale.com/a/abc123"
@patch("dropkit.main.subprocess.run")
def test_strips_trailing_punctuation(self, mock_run):
"""Test that trailing punctuation is stripped from URL."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"Visit: https://login.tailscale.com/a/abc123.\n",
)
url = run_tailscale_up("dropkit.test")
assert url == "https://login.tailscale.com/a/abc123"
@patch("dropkit.main.subprocess.run")
def test_no_url_in_output(self, mock_run):
"""Test when no URL is found in output."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"Some other output without URL",
)
url = run_tailscale_up("dropkit.test")
assert url is None
@patch("dropkit.main.subprocess.run")
def test_non_tailscale_url_ignored(self, mock_run):
"""Test that non-tailscale URLs are ignored."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"Visit https://example.com/login for more info",
)
url = run_tailscale_up("dropkit.test")
assert url is None
@patch("dropkit.main.subprocess.run")
def test_ssh_timeout(self, mock_run):
"""Test when SSH connection times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 30)
url = run_tailscale_up("dropkit.test")
assert url is None
class TestLockDownToTailscale:
"""Tests for lock_down_to_tailscale function."""
@patch("dropkit.main.subprocess.run")
def test_all_commands_succeed(self, mock_run):
"""Test when all UFW commands succeed."""
mock_run.return_value = MagicMock(returncode=0, stderr=b"")
result = lock_down_to_tailscale("dropkit.test")
assert result is True
# Should have called 5 commands
assert mock_run.call_count == 5
@patch("dropkit.main.subprocess.run")
def test_first_command_fails(self, mock_run):
"""Test when first UFW command fails."""
mock_run.return_value = MagicMock(returncode=1, stderr=b"ufw error")
result = lock_down_to_tailscale("dropkit.test")
assert result is False
# Should stop after first failure
assert mock_run.call_count == 1
@patch("dropkit.main.subprocess.run")
def test_middle_command_fails(self, mock_run):
"""Test when a middle command fails."""
# First two succeed, third fails
mock_run.side_effect = [
MagicMock(returncode=0, stderr=b""),
MagicMock(returncode=0, stderr=b""),
MagicMock(returncode=1, stderr=b"deny error"),
]
result = lock_down_to_tailscale("dropkit.test")
assert result is False
assert mock_run.call_count == 3
@patch("dropkit.main.subprocess.run")
def test_ssh_timeout(self, mock_run):
"""Test when SSH connection times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 30)
result = lock_down_to_tailscale("dropkit.test")
assert result is False
class TestVerifyTailscaleSsh:
"""Tests for verify_tailscale_ssh function."""
@patch("dropkit.main.subprocess.run")
def test_ssh_success(self, mock_run):
"""Test when SSH via Tailscale works."""
mock_run.return_value = MagicMock(returncode=0)
result = verify_tailscale_ssh("100.64.1.1", "testuser", "~/.ssh/id_ed25519")
assert result is True
@patch("dropkit.main.subprocess.run")
def test_ssh_failure(self, mock_run):
"""Test when SSH via Tailscale fails."""
mock_run.return_value = MagicMock(returncode=255)
result = verify_tailscale_ssh("100.64.1.1", "testuser", "~/.ssh/id_ed25519")
assert result is False
@patch("dropkit.main.subprocess.run")
def test_ssh_timeout(self, mock_run):
"""Test when SSH times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 15)
result = verify_tailscale_ssh("100.64.1.1", "testuser", "~/.ssh/id_ed25519")
assert result is False
class TestTailscaleConfig:
"""Tests for TailscaleConfig Pydantic model."""
def test_defaults(self):
"""Test default values."""
config = TailscaleConfig()
assert config.enabled is True
assert config.lock_down_firewall is True
assert config.auth_timeout == 300
def test_custom_values(self):
"""Test custom values."""
config = TailscaleConfig(enabled=False, lock_down_firewall=False, auth_timeout=600)
assert config.enabled is False
assert config.lock_down_firewall is False
assert config.auth_timeout == 600
def test_auth_timeout_minimum(self):
"""Test that auth_timeout has minimum validation."""
with pytest.raises(ValueError):
TailscaleConfig(auth_timeout=10) # Below minimum of 30
def test_auth_timeout_at_minimum(self):
"""Test auth_timeout at minimum value."""
config = TailscaleConfig(auth_timeout=30)
assert config.auth_timeout == 30
class TestCheckTailscaleInstalled:
"""Tests for check_tailscale_installed function."""
@patch("dropkit.main.subprocess.run")
def test_tailscale_installed(self, mock_run):
"""Test when Tailscale is installed."""
mock_run.return_value = MagicMock(returncode=0)
assert check_tailscale_installed("dropkit.test") is True
@patch("dropkit.main.subprocess.run")
def test_tailscale_not_installed(self, mock_run):
"""Test when Tailscale is not installed."""
mock_run.return_value = MagicMock(returncode=1)
assert check_tailscale_installed("dropkit.test") is False
@patch("dropkit.main.subprocess.run")
def test_ssh_timeout(self, mock_run):
"""Test when SSH connection times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 15)
assert check_tailscale_installed("dropkit.test") is False
@patch("dropkit.main.subprocess.run")
def test_ssh_connection_failed(self, mock_run):
"""Test when SSH connection fails."""
import subprocess
mock_run.side_effect = subprocess.SubprocessError("Connection refused")
assert check_tailscale_installed("dropkit.test") is False
class TestInstallTailscaleOnDroplet:
"""Tests for install_tailscale_on_droplet function."""
@patch("dropkit.main.subprocess.run")
def test_install_success(self, mock_run):
"""Test successful Tailscale installation."""
mock_run.return_value = MagicMock(
returncode=0,
stdout=b"Installation complete!",
)
assert install_tailscale_on_droplet("dropkit.test") is True
@patch("dropkit.main.subprocess.run")
def test_install_failure(self, mock_run):
"""Test failed Tailscale installation."""
mock_run.return_value = MagicMock(
returncode=1,
stdout=b"Error: curl failed",
)
assert install_tailscale_on_droplet("dropkit.test") is False
@patch("dropkit.main.subprocess.run")
def test_install_timeout(self, mock_run):
"""Test when installation times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 120)
assert install_tailscale_on_droplet("dropkit.test") is False
@patch("dropkit.main.subprocess.run")
def test_ssh_connection_failed(self, mock_run):
"""Test when SSH connection fails during install."""
import subprocess
mock_run.side_effect = subprocess.SubprocessError("Connection refused")
assert install_tailscale_on_droplet("dropkit.test") is False
def create_mock_config(lock_down_firewall: bool = True, auth_timeout: int = 300) -> MagicMock:
"""Create a mock DropkitConfig for testing."""
config = MagicMock()
config.tailscale = MagicMock()
config.tailscale.lock_down_firewall = lock_down_firewall
config.tailscale.auth_timeout = auth_timeout
config.ssh = MagicMock()
config.ssh.config_path = "~/.ssh/config"
config.ssh.identity_file = "~/.ssh/id_ed25519"
return config
class TestSetupTailscale:
"""Tests for setup_tailscale function."""
@patch("dropkit.main.verify_tailscale_ssh")
@patch("dropkit.main.lock_down_to_tailscale")
@patch("dropkit.main.check_local_tailscale")
@patch("dropkit.main.add_ssh_host")
@patch("dropkit.main.wait_for_tailscale_ip")
@patch("dropkit.main.run_tailscale_up")
def test_already_authenticated_succeeds(
self,
mock_tailscale_up,
mock_wait_ip,
mock_add_ssh,
mock_check_local,
mock_lockdown,
mock_verify,
):
"""Test setup succeeds when Tailscale is already authenticated (no auth URL)."""
# No auth URL returned (already authenticated)
mock_tailscale_up.return_value = None
# But Tailscale IP is available
mock_wait_ip.return_value = "100.64.1.1"
mock_check_local.return_value = True
mock_lockdown.return_value = True
mock_verify.return_value = True
config = create_mock_config()
result = setup_tailscale("dropkit.test", "testuser", config)
assert result == "100.64.1.1"
# Should have called wait_for_tailscale_ip with short timeout
mock_wait_ip.assert_called_once_with(
"dropkit.test", timeout=10, poll_interval=2, verbose=False
)
mock_add_ssh.assert_called_once()
mock_lockdown.assert_called_once()
@patch("dropkit.main.wait_for_tailscale_ip")
@patch("dropkit.main.run_tailscale_up")
def test_no_auth_url_and_not_connected_fails(
self,
mock_tailscale_up,
mock_wait_ip,
):
"""Test setup fails when no auth URL and Tailscale not connected."""
# No auth URL returned
mock_tailscale_up.return_value = None
# And no Tailscale IP available
mock_wait_ip.return_value = None
config = create_mock_config()
result = setup_tailscale("dropkit.test", "testuser", config)
assert result is None
@patch("dropkit.main.verify_tailscale_ssh")
@patch("dropkit.main.lock_down_to_tailscale")
@patch("dropkit.main.check_local_tailscale")
@patch("dropkit.main.add_ssh_host")
@patch("dropkit.main.wait_for_tailscale_ip")
@patch("dropkit.main.run_tailscale_up")
def test_normal_auth_flow_succeeds(
self,
mock_tailscale_up,
mock_wait_ip,
mock_add_ssh,
mock_check_local,
mock_lockdown,
mock_verify,
):
"""Test normal flow with auth URL succeeds."""
# Auth URL returned (needs authentication)
mock_tailscale_up.return_value = "https://login.tailscale.com/a/abc123"
# Tailscale IP available after authentication
mock_wait_ip.return_value = "100.64.1.1"
mock_check_local.return_value = True
mock_lockdown.return_value = True
mock_verify.return_value = True
config = create_mock_config()
result = setup_tailscale("dropkit.test", "testuser", config)
assert result == "100.64.1.1"
# Should have called wait_for_tailscale_ip with full auth_timeout
mock_wait_ip.assert_called_once_with("dropkit.test", timeout=300, verbose=False)
@patch("dropkit.main.wait_for_tailscale_ip")
@patch("dropkit.main.run_tailscale_up")
def test_auth_timeout_fails(
self,
mock_tailscale_up,
mock_wait_ip,
):
"""Test setup fails when authentication times out."""
# Auth URL returned
mock_tailscale_up.return_value = "https://login.tailscale.com/a/abc123"
# But no IP received (user didn't authenticate in time)
mock_wait_ip.return_value = None
config = create_mock_config()
result = setup_tailscale("dropkit.test", "testuser", config)
assert result is None
@patch("dropkit.main.check_local_tailscale")
@patch("dropkit.main.add_ssh_host")
@patch("dropkit.main.wait_for_tailscale_ip")
@patch("dropkit.main.run_tailscale_up")
def test_skips_lockdown_when_disabled(
self,
mock_tailscale_up,
mock_wait_ip,
mock_add_ssh,
mock_check_local,
):
"""Test firewall lockdown is skipped when disabled in config."""
mock_tailscale_up.return_value = None
mock_wait_ip.return_value = "100.64.1.1"
config = create_mock_config(lock_down_firewall=False)
result = setup_tailscale("dropkit.test", "testuser", config)
assert result == "100.64.1.1"
mock_check_local.assert_not_called() # Lockdown logic not entered
class TestTailscaleLogout:
"""Tests for tailscale_logout function."""
@patch("dropkit.main.subprocess.run")
def test_logout_success(self, mock_run):
"""Test successful Tailscale logout."""
mock_run.return_value = MagicMock(returncode=0, stderr=b"")
result = tailscale_logout("dropkit.test")
assert result is True
# Verify correct SSH command was called
mock_run.assert_called_once()
call_args = mock_run.call_args[0][0]
assert "ssh" in call_args
assert "dropkit.test" in call_args
assert "sudo tailscale logout" in call_args
@patch("dropkit.main.subprocess.run")
def test_logout_command_fails(self, mock_run):
"""Test when tailscale logout command fails."""
mock_run.return_value = MagicMock(returncode=1, stderr=b"logout failed")
result = tailscale_logout("dropkit.test")
assert result is False
@patch("dropkit.main.subprocess.run")
def test_ssh_connection_fails(self, mock_run):
"""Test when SSH connection fails."""
import subprocess
mock_run.side_effect = subprocess.SubprocessError("Connection refused")
result = tailscale_logout("dropkit.test")
assert result is False
@patch("dropkit.main.subprocess.run")
def test_ssh_timeout(self, mock_run):
"""Test when SSH connection times out."""
import subprocess
mock_run.side_effect = subprocess.TimeoutExpired("ssh", 30)
result = tailscale_logout("dropkit.test")
assert result is False
tests/test_version_check.py
+313
View File
@@ -0,0 +1,313 @@
"""Tests for version_check module."""
import json
import subprocess
import time
from unittest.mock import Mock, patch
from dropkit.version_check import (
commits_differ,
extract_commit_from_version,
get_last_check_file,
get_latest_git_commit,
should_check_version,
update_last_check_time,
)
class TestGetLastCheckFile:
"""Tests for get_last_check_file function."""
def test_returns_correct_path(self, tmp_path, monkeypatch):
"""Test that get_last_check_file returns correct path."""
# Mock Config.get_config_dir to return tmp_path
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
result = get_last_check_file()
assert result == tmp_path / ".last_version_check"
assert result.name == ".last_version_check"
class TestShouldCheckVersion:
"""Tests for should_check_version function."""
def test_no_file_returns_true(self, tmp_path, monkeypatch):
"""Test returns True when check file doesn't exist."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
result = should_check_version()
assert result is True
def test_expired_check_returns_true(self, tmp_path, monkeypatch):
"""Test returns True when more than 24 hours have passed."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
# Create check file with old timestamp (25 hours ago)
check_file = tmp_path / ".last_version_check"
old_timestamp = time.time() - (25 * 3600) # 25 hours ago
check_file.write_text(json.dumps({"timestamp": old_timestamp}))
result = should_check_version()
assert result is True
def test_not_expired_returns_false(self, tmp_path, monkeypatch):
"""Test returns False when less than 24 hours have passed."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
# Create check file with recent timestamp (1 hour ago)
check_file = tmp_path / ".last_version_check"
recent_timestamp = time.time() - 3600 # 1 hour ago
check_file.write_text(json.dumps({"timestamp": recent_timestamp}))
result = should_check_version()
assert result is False
def test_corrupted_file_returns_true(self, tmp_path, monkeypatch):
"""Test returns True when file contains invalid JSON."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
# Create check file with invalid JSON
check_file = tmp_path / ".last_version_check"
check_file.write_text("not valid json {{{")
result = should_check_version()
assert result is True
def test_missing_timestamp_returns_true(self, tmp_path, monkeypatch):
"""Test returns True when file missing timestamp key."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
# Create check file without timestamp
check_file = tmp_path / ".last_version_check"
check_file.write_text(json.dumps({"other_key": "value"}))
result = should_check_version()
# Should return True because timestamp defaults to 0, making it very old
assert result is True
class TestUpdateLastCheckTime:
"""Tests for update_last_check_time function."""
def test_creates_file_with_correct_structure(self, tmp_path, monkeypatch):
"""Test creates file with timestamp and current_version."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
before_time = time.time()
update_last_check_time()
after_time = time.time()
check_file = tmp_path / ".last_version_check"
assert check_file.exists()
data = json.loads(check_file.read_text())
assert "timestamp" in data
assert "current_version" in data
# Verify timestamp is recent (within the test execution window)
assert before_time <= data["timestamp"] <= after_time
def test_creates_parent_directory(self, tmp_path, monkeypatch):
"""Test creates parent directory if it doesn't exist."""
from dropkit.config import Config
nested_path = tmp_path / "nested" / "config"
monkeypatch.setattr(Config, "get_config_dir", lambda: nested_path)
update_last_check_time()
check_file = nested_path / ".last_version_check"
assert check_file.exists()
assert check_file.parent == nested_path
def test_silent_failure_on_permission_error(self, tmp_path, monkeypatch):
"""Test silently fails if can't write file."""
from dropkit.config import Config
monkeypatch.setattr(Config, "get_config_dir", lambda: tmp_path)
# Mock open to raise OSError
with patch("builtins.open", side_effect=OSError("Permission denied")):
# Should not raise exception
update_last_check_time()
class TestExtractCommitFromVersion:
"""Tests for extract_commit_from_version function."""
def test_format_git_dot(self):
"""Test extraction from '0.1.0+git.abc1234' format."""
result = extract_commit_from_version("0.1.0+git.abc1234")
assert result == "abc1234"
def test_format_plus_only(self):
"""Test extraction from '0.1.0+abc1234' format."""
result = extract_commit_from_version("0.1.0+abc1234")
assert result == "abc1234"
def test_truncates_long_hash(self):
"""Test truncates hash to 7 characters."""
result = extract_commit_from_version("0.1.0+git.abc1234567890")
assert result == "abc1234"
assert len(result) == 7
def test_dev_version_returns_none(self):
"""Test returns None for 'dev' version."""
result = extract_commit_from_version("dev")
assert result is None
def test_no_commit_returns_none(self):
"""Test returns None for version without commit."""
result = extract_commit_from_version("0.1.0")
assert result is None
def test_empty_string_returns_none(self):
"""Test returns None for empty string."""
result = extract_commit_from_version("")
assert result is None
def test_format_with_dots_in_suffix(self):
"""Test extraction when suffix contains dots."""
result = extract_commit_from_version("0.1.0+dev.123.abc1234")
# Should extract the last part after splitting by '.'
assert result == "abc1234"
class TestCommitsDiffer:
"""Tests for commits_differ function."""
def test_same_commit_returns_false(self):
"""Test returns False when commits match."""
result = commits_differ("0.1.0+git.abc1234", "abc1234")
assert result is False
def test_different_commits_returns_true(self):
"""Test returns True when commits differ."""
result = commits_differ("0.1.0+git.abc1234", "xyz9876")
assert result is True
def test_no_current_commit_returns_false(self):
"""Test returns False when current version has no commit."""
result = commits_differ("0.1.0", "abc1234")
assert result is False
def test_case_insensitive_comparison(self):
"""Test comparison is case-insensitive."""
result1 = commits_differ("0.1.0+git.ABC1234", "abc1234")
result2 = commits_differ("0.1.0+git.abc1234", "ABC1234")
assert result1 is False
assert result2 is False
def test_dev_version_returns_false(self):
"""Test returns False for dev version."""
result = commits_differ("dev", "abc1234")
assert result is False
class TestGetLatestGitCommit:
"""Tests for get_latest_git_commit function."""
@patch("subprocess.run")
def test_success_returns_short_hash(self, mock_run):
"""Test successful fetch returns 7-char hash."""
# Mock successful git ls-remote
mock_result = Mock()
mock_result.returncode = 0
mock_result.stdout = "abc1234567890abcdef1234567890abcdef1234\tHEAD\n"
mock_run.return_value = mock_result
result = get_latest_git_commit()
assert result == "abc1234"
assert len(result) == 7
@patch("subprocess.run")
def test_nonzero_returncode_returns_none(self, mock_run):
"""Test non-zero return code returns None."""
mock_result = Mock()
mock_result.returncode = 1
mock_result.stdout = ""
mock_run.return_value = mock_result
result = get_latest_git_commit()
assert result is None
@patch("subprocess.run")
def test_empty_output_returns_none(self, mock_run):
"""Test empty stdout returns None."""
mock_result = Mock()
mock_result.returncode = 0
mock_result.stdout = ""
mock_run.return_value = mock_result
result = get_latest_git_commit()
assert result is None
@patch("subprocess.run")
def test_timeout_returns_none(self, mock_run):
"""Test timeout returns None."""
mock_run.side_effect = subprocess.TimeoutExpired(cmd="git", timeout=5)
result = get_latest_git_commit()
assert result is None
@patch("subprocess.run")
def test_subprocess_error_returns_none(self, mock_run):
"""Test subprocess error returns None."""
mock_run.side_effect = subprocess.SubprocessError("Command failed")
result = get_latest_git_commit()
assert result is None
@patch("subprocess.run")
def test_os_error_returns_none(self, mock_run):
"""Test OS error returns None."""
mock_run.side_effect = OSError("Git not found")
result = get_latest_git_commit()
assert result is None
@patch("subprocess.run")
def test_calls_git_with_correct_args(self, mock_run):
"""Test calls git ls-remote with correct arguments."""
mock_result = Mock()
mock_result.returncode = 0
mock_result.stdout = "abc1234567890\tHEAD\n"
mock_run.return_value = mock_result
get_latest_git_commit()
mock_run.assert_called_once()
call_args = mock_run.call_args[0][0]
assert call_args == [
"git",
"ls-remote",
"https://github.com/trailofbits/dropkit.git",
"HEAD",
]
uv.lock
Generated
+1166
View File
File diff suppressed because it is too large Load Diff