From 9c51cabdfa586e9e70f709a98a03600f0fa6344c Mon Sep 17 00:00:00 2001
From: lza_menace <lza_menace@protonmail.com>
Date: Sun, 27 Sep 2020 20:52:58 -0700
Subject: [PATCH] add wallet cleanup functionality

---
 wowstash/blueprints/wallet/routes.py |  1 +
 wowstash/library/docker.py           | 19 ++++++++++++++++---
 wowstash/models.py                   |  2 ++
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/wowstash/blueprints/wallet/routes.py b/wowstash/blueprints/wallet/routes.py
index ed3c1a8..a96c3bb 100644
--- a/wowstash/blueprints/wallet/routes.py
+++ b/wowstash/blueprints/wallet/routes.py
@@ -84,6 +84,7 @@ def connect():
         current_user.wallet_connected = docker.container_exists(wallet)
         current_user.wallet_port = port
         current_user.wallet_container = wallet
+        current_user.wallet_start = datetime.utcnow()
         db.session.commit()
 
     return 'ok'
diff --git a/wowstash/library/docker.py b/wowstash/library/docker.py
index 2355cd9..fc381ba 100644
--- a/wowstash/library/docker.py
+++ b/wowstash/library/docker.py
@@ -3,6 +3,8 @@ from docker.errors import NotFound, NullResource, APIError
 from socket import socket
 from os.path import expanduser
 from secrets import token_urlsafe
+from datetime import datetime, timedelta
+from time import sleep
 from wowstash import config
 from wowstash.models import User
 from wowstash.factory import db
@@ -140,9 +142,20 @@ class Docker(object):
     def cleanup(self):
         users = User.query.all()
         for u in users:
-            if u.wallet_container:
-                if not self.container_exists(u.wallet_container):
-                    u.clear_wallet_data()
+            # Delete inactive wallet sessions
+            if u.wallet_start:
+                session_lifetime = getattr(config, 'PERMANENT_SESSION_LIFETIME', 3600)
+                expiration_time = u.wallet_start + timedelta(seconds=session_lifetime)
+                now = datetime.utcnow()
+                time_diff = expiration_time - now
+                if time_diff.total_seconds() <= 0:
+                    print(f'[+] Found expired container for {u}. killing it')
+                    self.stop_container(u.wallet_container)
+                    sleep(2)
+            # Remove wallet db data if not running but it's in db
+            if u.wallet_container and not self.container_exists(u.wallet_container):
+                print(f'[+] Found stale data for {u}')
+                u.clear_wallet_data()
 
 
 docker = Docker()
diff --git a/wowstash/models.py b/wowstash/models.py
index ac634b8..2d87caf 100644
--- a/wowstash/models.py
+++ b/wowstash/models.py
@@ -19,6 +19,7 @@ class User(db.Model):
     wallet_connected = db.Column(db.Boolean, default=False)
     wallet_port = db.Column(db.Integer, nullable=True)
     wallet_container = db.Column(db.String(30), nullable=True)
+    wallet_start = db.Column(db.DateTime, nullable=True)
 
     @property
     def is_authenticated(self):
@@ -43,6 +44,7 @@ class User(db.Model):
         self.wallet_connected = False
         self.wallet_port = None
         self.wallet_container = None
+        self.wallet_start = None
         if reset_password:
             self.wallet_password = None
         if reset_wallet: