From cb46a0a6a5123e820de4bd6a797732770174def4 Mon Sep 17 00:00:00 2001
From: Graham Christensen <graham@grahamc.com>
Date: Tue, 24 Mar 2020 14:12:13 -0400
Subject: [PATCH] Add an s3 backend for kicks

---
 nixops/plugin.py     |  3 +-
 nixops/storage/s3.py | 82 ++++++++++++++++++++++++++++++++++++++++++++
 release.nix          |  1 +
 3 files changed, 85 insertions(+), 1 deletion(-)
 create mode 100644 nixops/storage/s3.py

diff --git a/nixops/plugin.py b/nixops/plugin.py
index 97481c43a..67ab3775d 100644
--- a/nixops/plugin.py
+++ b/nixops/plugin.py
@@ -3,8 +3,9 @@
 from nixops.storage import StorageBackend
 from nixops.storage.legacy import LegacyBackend
 from nixops.storage.memory import MemoryBackend
+from nixops.storage.s3 import S3Backend
 
 
 @nixops.plugins.hookimpl
 def register_backends() -> Dict[str, Type[StorageBackend]]:
-    return {"legacy": LegacyBackend, "memory": MemoryBackend}
+    return {"legacy": LegacyBackend, "s3": S3Backend, "memory": MemoryBackend}
diff --git a/nixops/storage/s3.py b/nixops/storage/s3.py
new file mode 100644
index 000000000..b5b97deec
--- /dev/null
+++ b/nixops/storage/s3.py
@@ -0,0 +1,82 @@
+from __future__ import annotations
+from nixops.storage import StorageArgDescriptions, StorageArgValues
+import boto3
+from botocore.exceptions import ClientError
+import sys
+import os
+import typing
+from typing import Dict
+
+if typing.TYPE_CHECKING:
+    import nixops.statefile
+
+
+class S3Backend:
+    @staticmethod
+    def arguments() -> StorageArgDescriptions:
+        raise NotImplementedError
+
+    def __init__(self, args: StorageArgValues) -> None:
+        self.bucket = args["bucket"]
+        self.key = args["key"]
+        self.region = args["region"]
+        self.profile = args["profile"]
+        self.dynamodb_table = args["dynamodb_table"]
+        self.s3_endpoint = args.get("s3_endpoint")
+        self.kms_keyid = args.get("kms_keyid")
+        self.aws = boto3.Session(region_name=self.region, profile_name=self.profile)
+
+    # fetchToFile: acquire a lock and download the state file to
+    # the local disk. Note: no arguments will be passed over kwargs.
+    # Making it part of the type definition allows adding new
+    # arguments later.
+    def fetchToFile(self, path: str, **kwargs) -> None:
+        self.lock(path)
+        try:
+            with open(path, "wb") as f:
+                self.aws.client("s3").download_fileobj(self.bucket, self.key, f)
+            print("Fetched!")
+        except ClientError as e:
+            from pprint import pprint
+
+            pprint(e)
+            if e.response["Error"]["Code"] == "404":
+                self.aws.client("s3").put_object(
+                    Bucket=self.bucket, Key=self.key, Body=b"", **self.encargs()
+                )
+
+    def onOpen(self, sf: nixops.statefile.StateFile, **kwargs) -> None:
+        pass
+
+    # uploadFromFile: upload the new state file and release any locks
+    # Note: no arguments will be passed over kwargs. Making it part of
+    # the type definition allows adding new arguments later.
+    def uploadFromFile(self, path: str, **kwargs) -> None:
+        with open(path, "rb") as f:
+            self.aws.client("s3").upload_fileobj(
+                f, self.bucket, self.key, ExtraArgs=self.encargs()
+            )
+
+        self.unlock(path)
+
+    def s3(self) -> None:
+        self.aws.client("s3", endpoint_url=self.s3_endpoint)
+
+    def encargs(self) -> Dict[str, str]:
+        if self.kms_keyid is not None:
+            return {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": self.kms_keyid}
+        else:
+            return {}
+
+    def lock(self, path) -> None:
+        r = self.aws.client("dynamodb").put_item(
+            TableName=self.dynamodb_table,
+            Item={"LockID": {"S": f"{self.bucket}/{self.key}"},},
+            ConditionExpression="attribute_not_exists(LockID)",
+        )
+
+    def unlock(self, path: str) -> None:
+        self.aws.client("dynamodb").delete_item(
+            TableName=self.dynamodb_table,
+            Key={"LockID": {"S": f"{self.bucket}/{self.key}"},},
+        )
diff --git a/release.nix b/release.nix
index 9ea2e8ca9..cc2504a02 100644
--- a/release.nix
+++ b/release.nix
@@ -89,6 +89,7 @@ in rec {
           pythonPackages.prettytable
           pythonPackages.pluggy
           pythonPackages.typing-extensions
+          pythonPackages.boto3
         ] ++ pkgs.lib.traceValFn
            (x: "Using plugins: " + builtins.toJSON x)
            (map (d: d.build.${system}) (pluginSet allPlugins));