From b957200debf0cf85658fef3e0c1ab812156b38eb Mon Sep 17 00:00:00 2001
From: Fuxino <dfucini@gmail.com>
Date: Wed, 31 May 2023 19:30:31 +0200
Subject: [PATCH] Change missing hostkey policy

---
 simple_backup/simple_backup.py | 41 +++++++++++++++++++++++-----------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/simple_backup/simple_backup.py b/simple_backup/simple_backup.py
index 5fb203e..5c98c33 100755
--- a/simple_backup/simple_backup.py
+++ b/simple_backup/simple_backup.py
@@ -14,6 +14,7 @@ Classes:
 # Import libraries
 import sys
 import os
+import warnings
 from functools import wraps
 from shutil import rmtree
 import shlex
@@ -31,6 +32,8 @@ from dotenv import load_dotenv
 import paramiko
 from paramiko import RSAKey, Ed25519Key, ECDSAKey, DSSKey
 
+warnings.filterwarnings('error')
+
 
 try:
     from systemd import journal
@@ -298,30 +301,42 @@ class Backup:
 
     def _ssh_connection(self):
         ssh = paramiko.SSHClient()
-        ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        ssh.load_system_host_keys()
+        ssh.set_missing_host_key_policy(paramiko.WarningPolicy())
 
-        agent = paramiko.Agent()
-        agent_keys = agent.get_keys()
+        try:
+            ssh.connect(self.host, username=self.username)
 
-        for key in agent_keys:
-            try:
-                ssh.connect(self.host, username=self.username, pkey=key)
-                return ssh
-            except paramiko.SSHException:
-                pass
+            return ssh
+        except UserWarning:
+            k = input(f'Unknown key for host {self.host}. Continue anyway? (Y/N) ')
+
+            if k[0].upper() == 'Y':
+                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+            else:
+                return None
+        except paramiko.SSHException:
+            pass
+
+        try:
+            ssh.connect(self.host, username=self.username)
+
+            return ssh
+        except paramiko.SSHException:
+            pass
 
         pkey = None
         password = None
 
         if self.ssh_keyfile is None:
-            logger.critical('Can\'t connect to the server. No key specified')
+            logger.critical('Can\'t connect to the server. No authentication method available')
 
             return None
 
         try:
             pkey = RSAKey.from_private_key_file(self.ssh_keyfile)
         except paramiko.PasswordRequiredException:
-            password = getpass()
+            password = getpass(f'Enter passwphrase for key \'{self.ssh_keyfile}\': ')
 
             try:
                 pkey = RSAKey.from_private_key_file(self.ssh_keyfile, password)
@@ -357,8 +372,8 @@ class Backup:
 
         try:
             ssh.connect(self.host, username=self.username, pkey=pkey)
-        except paramiko.SSHException as e:
-            logger.error(e)
+        except paramiko.SSHException:
+            logger.critical('SSH connection to server failed')
 
             return None