From 638a32b61a58c2eb03769ed9757e736228e4840b Mon Sep 17 00:00:00 2001
From: Raphael Roberts <raphael.roberts48@gmail.com>
Date: Wed, 26 May 2021 11:05:19 -0500
Subject: [PATCH] Fixed up main function and added prune feature

---
 openwrt_backup/__init__.py | 59 ++++++++++++++++++++++++++++++++++----
 openwrt_backup/prune.py    |  7 +++++
 2 files changed, 61 insertions(+), 5 deletions(-)
 create mode 100644 openwrt_backup/prune.py

diff --git a/openwrt_backup/__init__.py b/openwrt_backup/__init__.py
index 6c6499d..e1bd0db 100644
--- a/openwrt_backup/__init__.py
+++ b/openwrt_backup/__init__.py
@@ -1,13 +1,25 @@
 from argparse import ArgumentParser
 from pathlib import Path
 
+import toml
+
 from openwrt_backup.retrieve import make_backup_on_remote, retrieve_backup_on_remote
 from openwrt_backup.utils import (
     client_from_config,
     exec_remote,
     get_new_backup_filepath,
 )
-from openwrt_backup.config import CONFIG_DATA
+from openwrt_backup.config import CONFIG_DATA  # noqa
+from openwrt_backup.prune import get_all_but_last_n
+
+
+def do_prune(dest_path: Path, n: int, dry):
+    to_delete = get_all_but_last_n(dest_path, n)
+    for f in to_delete:
+        if dry:
+            print(f"Deleting: {f}")
+        else:
+            f.unlink()
 
 
 def do_backup(
@@ -33,16 +45,47 @@ def do_backup(
 
 
 def main():
+    global CONFIG_DATA
+
     parser = ArgumentParser()
-    parser.add_argument("-d", "--dry", action="store_true", help="Don't do anything")
-    parser.add_argument("--dump-config", action="store_true", help="Show config")
+
     parser.add_argument(
-        "-c" "--config", help="config path for backup script", default=None
+        "-d",
+        "--dry",
+        action="store_true",
+        help="Don't do anything but show what would happen",
     )
     parser.add_argument(
-        "name", help="hostname in ssh_config file for openwrt device we need to backup"
+        "-c", "--config", help="config path for backup script", default=None
+    )
+
+    subparsers = parser.add_subparsers(dest="action", required=True)
+    p_run = subparsers.add_parser(
+        "run",
+        help=f"Make backup then prune all but the {CONFIG_DATA['backup']['keep']} most recent",
+    )
+    p_backup = subparsers.add_parser(
+        "backup", help="Same as run but do not delete any backups"
+    )
+    p_dump_config = subparsers.add_parser(  # noqa
+        "dump-config", help="Show config and exit"
+    )  # noqa
+    p_run.add_argument(
+        "name", help="Hostname in ssh_config file for openwrt device we need to backup"
+    )
+    p_backup.add_argument(
+        "name", help="Hostname in ssh_config file for openwrt device we need to backup"
     )
     args = parser.parse_args()
+
+    if args.config:
+        with open(args.config) as f:
+            CONFIG_DATA = toml.load(f)
+
+    if args.action == "dump-config":
+        print(toml.dumps(CONFIG_DATA))
+        exit(0)
+
     do_backup(
         args.name,
         Path(CONFIG_DATA["ssh"]["root"]).expanduser(),
@@ -51,3 +94,9 @@ def main():
         CONFIG_DATA["backup"]["time_format"],
         args.dry,
     )
+    if args.action != "backup":
+        do_prune(
+            Path(CONFIG_DATA["backup"]["path"]).expanduser(),
+            CONFIG_DATA["backup"]["keep"],
+            args.dry,
+        )
diff --git a/openwrt_backup/prune.py b/openwrt_backup/prune.py
new file mode 100644
index 0000000..20e3f4c
--- /dev/null
+++ b/openwrt_backup/prune.py
@@ -0,0 +1,7 @@
+from pathlib import Path
+
+
+def get_all_but_last_n(p: Path, n: int):
+    files = list(p.glob("*.tar.gz"))
+    sorted_files = sorted(files, key=lambda f: int(f.name.split("__")[0]), reverse=True)
+    return sorted_files[n:]