Add SSH Key Auth [Follow-up]

.env.example

@@ -8,7 +8,12 @@ conn.port=2222
 
 # exchange with the user for your target VM
 conn.username='bob'
+#To just use keyauth only, use '' with no space for conn.password 
+#Otherwise, insert the password for instance here
 conn.password='secret'
+#To just use username and password auth only, use '' with no space for conn.keyfilename
+#Otherwise, insert the filepath for the keyfile here (for example, '/home/bob/.ssh/sshkey.rsa')
+conn.keyfilename=''
 
 # which LLM model to use (can be anything openai supports, or if you use a custom llm.api_url, anything your api provides for the model parameter
 llm.model='gpt-3.5-turbo'

.env.example.aws

@@ -0,0 +1,23 @@
+llm.api_key='your-openai-key'
+log_db.connection_string='log_db.sqlite3'
+
+# exchange with the IP of your target VM
+conn.host='enter the public IP of AWS Instance'
+conn.hostname='DNS of AWS Instance '
+conn.port=22
+
+# user of target AWS Instance
+conn.username='bob'
+#To just use keyauth only, use '' with no space for conn.password 
+#Otherwise, insert the password for instance here
+conn.password=''
+#To just use username and password auth only, use '' with no space for conn.keyfilename
+#Otherwise, insert the filepath for the keyfile here (for example, '/home/bob/.ssh/awskey.pem')
+conn.keyfilename='/home/bob/.ssh/awskey.pem'
+
+# which LLM model to use (can be anything openai supports, or if you use a custom llm.api_url, anything your api provides for the model parameter
+llm.model='gpt-3.5-turbo'
+llm.context_size=16385
+
+# how many rounds should this thing go?
+max_turns = 20

README.md

@@ -166,9 +166,12 @@ $ source ./venv/bin/activate
 # install python requirements
 $ pip install -e .
 
-# copy default .env.example
+# copy default .env.example 
 $ cp .env.example .env
 
+# NOTE: if you are trying to use this with AWS or ssh-key only authentication, copy .env.example.aws
+$ cp .env.example.aws .env 
+
 # IMPORTANT: setup your OpenAI API key, the VM's IP and credentials within .env
 $ vi .env
 

src/hackingBuddyGPT/capabilities/ssh_test_credential.py

@@ -18,8 +18,8 @@ def describe(self) -> str:
     def get_name(self):
         return "test_credential"
 
-    def __call__(self, username: str, password: str) -> Tuple[str, bool]:
-        test_conn = self.conn.new_with(username=username, password=password)
+    def __call__(self, username: str, password: str, keyfilename: str) -> Tuple[str, bool]:
+        test_conn = self.conn.new_with(username=username, password=password, keyfilename=keyfilename)
         try:
             test_conn.init()
             user = test_conn.run("whoami")[0].strip("\n\r ")

src/hackingBuddyGPT/utils/ssh_connection/ssh_connection.py

@@ -14,25 +14,33 @@ class SSHConnection:
     hostname: str
     username: str
     password: str
+    keyfilename: str
     port: int = 22
 
     _conn: Connection = None
 
     def init(self):
         # create the SSH Connection
-        conn = Connection(
-            f"{self.username}@{self.host}:{self.port}",
-            connect_kwargs={"password": self.password, "look_for_keys": False, "allow_agent": False},
-        )
+        if self.keyfilename == '' or self.keyfilename == None:
+            conn = Connection(
+                f"{self.username}@{self.host}:{self.port}",
+                connect_kwargs={"password": self.password, "look_for_keys": False, "allow_agent": False},
+            )
+        else: 
+            conn = Connection(
+                f"{self.username}@{self.host}:{self.port}",
+                connect_kwargs={"password": self.password, "key_filename": self.keyfilename, "look_for_keys": False, "allow_agent": False},
+            )
         self._conn = conn
         self._conn.open()
 
-    def new_with(self, *, host=None, hostname=None, username=None, password=None, port=None) -> "SSHConnection":
+    def new_with(self, *, host=None, hostname=None, username=None, password=None, keyfilename=None, port=None) -> "SSHConnection":
         return SSHConnection(
             host=host or self.host,
             hostname=hostname or self.hostname,
             username=username or self.username,
             password=password or self.password,
+            keyfilename=keyfilename or self.keyfilename,
             port=port or self.port,
         )