Develop (#9)

* Refactor and cleanup
* Add more tests
* Add create and delete commands
* Add help output
* Update readme and other fixes

.dockerignore

@@ -0,0 +1,6 @@
+.venv
+venv
+.vagrant
+vagrant
+.mypy_cache
+.git

.gitignore

@@ -1,10 +1,10 @@
 .vscode
 .idea
 *.sw?
-*.img
 *.vdi
 stubs
 tmp
 venv
 dist
 __pycache__
+samconfig.toml

Makefile

@@ -22,7 +22,7 @@ docker/build:
 	docker build -f Dockerfile.mount -t dsnap-mount .
 
 docker/run:
-	docker run -it -v "${PWD}/${IMAGE}:/disks/${IMAGE}" -w /disks mount --ro -a "${IMAGE}" -m /dev/sda1:/
+	docker run -it -v "${PWD}/${IMAGE}:/disks/${IMAGE}" -w /disks dsnap-mount --ro -a "${IMAGE}" -m /dev/sda1:/
 
 test:
 	pytest ./tests

README.md

@@ -8,29 +8,62 @@ Utility for downloading EBS snapshots using the EBS Direct API's.
 
 ### PyPi
 
+```shell
+=======
 ```
 % pip install -U pip
 % pip install 'dsnap[cli]'
 ```
 
+## Command Reference
+
+```shell
+% dsnap --help
+Usage: dsnap [OPTIONS] COMMAND [ARGS]...
+
+  A utility for managing snapshots via the EBS Direct API.
+
+Options:
+  --region REGION                 Sets the AWS region.  [default: us-east-1]
+  --profile PROFILE               Shared credential profile to use.
+  --install-completion [bash|zsh|fish|powershell|pwsh]
+                                  Install completion for the specified shell.
+  --show-completion [bash|zsh|fish|powershell|pwsh]
+                                  Show completion for the specified shell, to
+                                  copy it or customize the installation.
+
+  --help                          Show this message and exit.
+
+Commands:
+  create  Create a snapshot for the given instances default device volume.
+  delete  Delete a given snapshot.
+  get     Download a snapshot for a given instance or snapshot ID.
+  init    Write out a Vagrantfile template to explore downloaded snapshots.
+  list    List snapshots in AWS.
+```
+
 ## Examples
 
+### Recording
+
+[![asciicast](https://asciinema.org/a/391559.svg)](https://asciinema.org/a/391559)
+
 ### Listing Snapshots
-```
-% dsnap --profile demo list
+```shell
+% dsnap list
            Id          |   Owner ID   |   State
 snap-0dbb0347f47e38b96   922105094392   completed
 ```
 
 ### Downloading a Snapshot
-```
-% dsnap --profile demo get snap-0dbb0347f47e38b96
+```shell
+% dsnap get snap-0dbb0347f47e38b96
 Output Path: /cwd/snap-0dbb0347f47e38b96.img
 ```
 
 If you don't specify a snapshot  you'll get a prompt to ask which one you want to download:
-```
-% python -m dsnap --profile chris get
+```shell
+% dsnap chris get
 0) i-01f0841393cd39f06 (ip-172-31-27-0.ec2.internal, vpc-04a91864355539a41, subnet-0e56cd55282fa9158)
 Select Instance: 0
 0) vol-0a1aab48b0bc3039d (/dev/sdb)
@@ -46,7 +79,7 @@ Cleaning up snapshot: snap-0543a8681adce0086
 ### Mounting in Vagrant
 This requires virtualbox to be installed. dsnap init will write a Vagrantfile to the current directory that can be used to mount a specific downloaded snapshot. Conversion to a VDI disk is handled in the Vagrantfile, it will look for the disk file specified in the IMAGE environment variable, convert it to a VDI using `VBoxManage convertdd`. The resulting VDI is destroyed when the Vagrant box is, however the original raw .img file will remain and can be reused as needed.
 
-```
+```shell
 % dsnap init
 % IMAGE=snap-0543a8681adce0086.img vagrant up
 % vagrant ssh
@@ -57,16 +90,16 @@ This requires virtualbox to be installed. dsnap init will write a Vagrantfile to
 This uses libguestfs to work directly with the downloaded img file.
 
 #### Build Docker Container
-```
-git clone https://github.com/RhinoSecurityLabs/dsnap.git
-cd dsnap
-make docker/build
+```shell
+% git clone https://github.com/RhinoSecurityLabs/dsnap.git
+% cd dsnap
+% make docker/build
 ```
 
 #### Run Guestfish Shell
 
-```
-IMAGE=snap-0dbb0347f47e38b96.img make docker/run
+```shell
+% IMAGE=snap-0dbb0347f47e38b96.img make docker/run
 ```
 
 This will take a second to start up. After it drops you into the shell you should be able to run commands like ls, cd, cat. However worth noting they don't always behave exactly like they do in a normal shell.
@@ -75,7 +108,7 @@ The output will give you the basics of how to use the guestfish shell. For a ful
 
 Below is an example of starting the shell and printing the contents of /etc/os-release.
 
-```
+```shell
 % IMAGE=snap-0dbb0347f47e38b96.img make docker/run
 docker run -it -v "/cwd/dsnap/snap-0dbb0347f47e38b96.img:/disks/snap-0dbb0347f47e38b96.img" -w /disks mount --ro -a "snap-0dbb0347f47e38b96.img" -m /dev/sda1:/
 
@@ -104,26 +137,26 @@ HOME_URL="https://amazonlinux.com/"
 For CLI development make sure you include the `cli` extra shown below. You'll also want to invoke the package by using python's `-m` (shown below) for testing local changes, the dnsap binary installed to the environment will only update when you run pip install.
 
 ### Setup
-```
-git clone https://github.com/RhinoSecurityLabs/dsnap.git
-cd dsnap
-python3 -m venv venv
-. venv/bin/activate
-python -m pip install '.[cli]'
+```shell
+% git clone https://github.com/RhinoSecurityLabs/dsnap.git
+% cd dsnap
+% python3 -m venv venv
+% . venv/bin/activate
+% python -m pip install '.[cli]'
 ```
 
 ### Running With Local Changes
-```
-python -m dsnap --help
+```shell
+% python -m dsnap --help
 ```
 
 ### Linting and Type Checking
-```
-make lint
+```shell
+% make lint
 ```
 
 ### Testing
-```
-make test
+```shell
+% make test
 ```
 

dsnap/__main__.py

@@ -6,6 +6,6 @@
 logging.basicConfig(format='[%(levelname)s] %(message)s', level=logging.WARNING)
 
 try:
-    app()
+    app(prog_name="dsnap")
 except (NoCredentialsError, NoRegionError) as e:
     logging.error(e.args[0])

dsnap/main.py

@@ -1,20 +1,18 @@
-import logging
-
-import sys
 from pathlib import Path
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Optional, List
 
 import boto3
-from typer import Argument, Option, Typer
+import typer
+from typer import Option, Typer
 
 from dsnap import utils
-from dsnap.snapshot import Snapshot
-from dsnap.prompt import snap_id_from_input
+from dsnap.prompt import snap_from_input, download_snap_id, snaps_from_input, take_snapshot, vol_from_id
+from dsnap.utils import fatal
 
 if TYPE_CHECKING:
     from mypy_boto3_ec2 import service_resource as r
 
-app = Typer()
+app = Typer(name="dsnap", help="Utility for downloading EBS snapshots using the EBS Direct API's.")
 
 sess: boto3.session.Session = boto3.session.Session()
 
@@ -23,54 +21,132 @@
 
 
 @app.callback()
-def session(region: str = Option(default='us-east-1'), profile: str = Option(default=None)):
-    """This is function set's up various global settings.
-
-    It is called by Typer before any of the other commands run due to the @app.callback decorator.
-    """
+def session(
+        region: str = Option(default='us-east-1', help="Sets the AWS region.", metavar="REGION"),
+        profile: str = Option(default=None, help="Shared credential profile to use.", metavar="PROFILE")
+):
     global sess, ec2
     sess = boto3.session.Session(region_name=region, profile_name=profile)
     ec2 = sess.resource('ec2')
 
 
 @app.command()
-def init(out_dir: Path = Path('.'), force: bool = False):
+def init(
+        out_dir: Path = typer.Option(Path('.'), help='Output directory to write Vagrantfile'),
+        force: bool = typer.Option(False, help='Overwrite any existing Vagrantfile.')
+):
+    """
+    Write out a Vagrantfile template to explore downloaded snapshots.
+
+    If --out-dir is used the given directory will be used instead.
+    If --force is used we will overwrite an already present Vagrantfile
+
+    To use the outputed Vagrantfile set the IMAGE environment to the path of the snapshot you want to mount and run vagrant
+    up. For example:
+
+    % dsnap init
+    % IMAGE=snap-0543a8681adce0086.img vagrant up
+    % vagrant ssh
+    """
     output = utils.init_vagrant(out_dir, force)
     if output:
-        print(f"Wrote Vagrantfile to {output}")
+        print(f"Wrote Vagrantfile to {typer.style('./'+str(output), bold=True)}")
     else:
         print(f"Vagrantfile already exists at {output}, use the --force to overwrite.")
 
 
 @app.command("list")
-def list_snapshots():
-    print("           Id          |   Owneer ID   | Description")
-    ec2: 'r.EC2ServiceResource' = sess.resource('ec2')
-    for snap in ec2.snapshots.filter(OwnerIds=['self']).all():
-        print(f"{snap.id}   {snap.owner_id}   {snap.description}")
+def list_snapshots(
+        instance_id: str = typer.Argument(None, help='Optional instance ID to limit listed snapshots to.'),
+        devices: List[str] = typer.Option(['/dev/sda', '/dev/xvda'], help='Optional device name to limit snapshots to.'),
+):
+    """
+    List snapshots in AWS.
+
+    If --instance-id is used then snapshots will be limited to that instances default device attachments.
+    If --devices is used alongside --instance-id then listed snapshots are for that instances given devices, by default this
+    is /dev/sda and /dev/xvda.
+    """
+    print(typer.style("           Id          |   Owneer ID   | Description   ", underline=True))
+    try:
+        for snap in snaps_from_input(sess, instance_id, devices):
+            print(f"{typer.style(snap.id, bold=True)}   {snap.owner_id}   {snap.description}")
+    except UserWarning as e:
+        fatal(*e.args)
 
 
 @app.command()
-def get(id: str = Argument(None), output: Optional[Path] = None):
-    # snap_id will be None in cases of invalid argument or no snapshot was selected
+def get(
+        # We use the filename to determine the snapshot id so we can only use directories for the output option.
+        output: Path = typer.Option(
+            Path('.'),
+            file_okay=False,
+            dir_okay=True,
+            help='If specified output the snapshot to the given directory, the name however is always the snapshot id.',
+        ),
+        force: bool = typer.Option(False, help='If specified and the snapshot already exists then overwrite it.'),
+        ids: Optional[List[str]] = typer.Argument(default=None, help='The remote snapshot ID to fetch.')
+):
+    """
+    Download a snapshot for a given instance or snapshot ID.
+
+    If no Argument is passed then you'll be prompted to select an instance, volume and snapshot to download. If no snapshot
+    exists, you can optionally create a temporary one.
+
+    If an instance ID is passed a snapshot for that instance will be downloaded, if more then one exists you'll be prompted
+    to select a one.
+
+    If a snapshot ID is passed that snapshot will be downloaded and you will not be prompted for any additional info.
+    """
     try:
-        snap_id = snap_id_from_input(sess, id)
-    except UserWarning as e:
-        print(*e.args, '\nExiting...')
-        sys.exit(1)
+        if not ids:
+            snap = snap_from_input(sess, ids)
+            download_snap_id(sess, force, output, snap.id)
+        else:
+            for id in ids:
+                snap = snap_from_input(sess, id)
+                download_snap_id(sess, force, output, snap.id)
+    except (UserWarning, FileExistsError) as e:
+        fatal(*e.args)
+
 
+@app.command()
+def create(ids: List[str] = typer.Argument(
+    None,
+    help='One or more ID\'s of a instance or volume to create a snapshot for. To avoid being prompted use an explict volume ID'
+         ' rather then an instance ID.'
+)):
+    """
+    Create a snapshot for the given instances default device volume.
+
+    The passed argument should be an instance ID, where a snapshot will be created from the default device volume, either
+    /dev/sda or /dev/xvda.
+    """
     try:
-        logging.info(f"Selected snapshot with id {snap_id}")
-        snap = Snapshot(snap_id, boto3_session=sess)
-        path = output and output.absolute().as_posix()
-        snap.download(path or f"{snap.snapshot_id}.img")
+        if not ids:
+            fatal("must pass at least one instance or volume id as an argument")
+        for i in ids:
+            vol = vol_from_id(sess, i)
+            s = take_snapshot(vol)
+            print(f"Created snapshot {typer.style(s.id, bold=True)} from instance " f"{typer.style(i, bold=True)}")
+
     except UserWarning as e:
-        print(*e.args)
-        sys.exit(2)
-    except Exception as e:
-        resp = getattr(e, 'response', None)
-        if resp and resp['Error']['Message']:
-            print(resp['Error']['Message'])
-            sys.exit(1)
-        else:
-            raise e
+        fatal(*e.args)
+
+
+@app.command()
+def delete(ids: List[str] = typer.Argument(None, help='One or more ID\'s of snapshots to delete')):
+    """
+    Delete a given snapshot.
+
+    The passed argument should be a snapshot ID to delete.
+    """
+    if not ids:
+        fatal("must pass at least one instance id as an argument")
+    for i in ids:
+        try:
+            s = ec2.Snapshot(i)
+            s.delete()
+            print(f"Deleted snapshot {typer.style(s.id, bold=True)}")
+        except UserWarning as e:
+            fatal(*e.args)