I am doing a project to ollama help me to maintain a Linux infrastructure, and need the IA to connect to a given ssh server with given credentials and execute a command and bring the output of command to the chat.
I already did this code, but it's not working, somethings (depending how I ask the prompt) I can see the ssh connection on the server, but no return of output.
This is the tool code I am using:
"""
title: ssh connection
description
"""
import os
import requests
from pydantic import BaseModel, Field
import paramiko
class Tools:
class Valves(BaseModel):
SSH_HOST: str = Field(
default="111.111.111.111", description="Adress to connect to ssh"
)
SSH_USERNAME: str = Field(
default="adminuser", description="Username to use on ssh server"
)
SSH_PASSWORD: str = Field(
default="thePassword123!", description="Password to connect on ssh server"
)
pass
class UserValves(BaseModel):
SSH_COMMAND: str = Field(
default="ifconfig",
description="user given command trought the chat",
)
pass
def __init__(self):
self.citation = True
self.Valves = self.Valves()
pass
# pass
def connect_ssh(
self, ssh_ip: str, ssh_username: str, ssh_password: str, ssh_command: str
) -> str:
"""
create a ssh connection using the data the user entered on the prompt,
This is the way you can connect to a ssh server, get the ssh_ip, ssh_username, ssh_password and ssh_command from the user prompt em use it to call the function paramiko.
use the bellow function to connect ssh servers and do what the user are asking
:param: ssh_command Command to be executed on ssh connection
:param: ssh_username username to log in in to ssh server
:param: ssh_password password to use on log in to ssh server
:param: ssh_host hostname or ip to open ssh connection
:output: A string containing the output got after command issued, and will be returned to user.
"""
"""
check if the informations for connect ssh is fullfilled
"""
if not self.Valves.SSH_IP:
return "Let the user know the ip to ssh connection was not provided. User will tell the ssh_host trought chat" ###
if not self.Valves.SSH_USERNAME:
return "Let the user know the username to ssh connection was not provided. User will tell the ssh_username trought chat"
if not self.Valves.SSH_PASSWORD:
return "Let the user know the password was not provided. User will tell the ssh_password trought chat"
if not self.UserValves.SSH_COMMAND:
return "Let the user know the ssh command was not provided. User will tell the ssh_command trought chat"
"""
use this function to connect to a ssh server, using parameters given by the user
"""
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(
self.UserValves.SSH_IP,
username=self.Valves.SSH_USERNAME,
password=self.Valves.SSH_PASSWORD,
)
stdin, stdout, stderr = ssh.exec_command(self.UserValves.SSH_COMMAND)
"""
get the output of the command and return to the chat.
"""
if stdout:
output += stdout.read().decode()
if stderr:
output += stderr.read().decode()
return "say that the command was executed, and you are waiting the output to tell on chat"
ssh.close()
return output
I think you are on the right track, here are some issues I could see
Issues:
You have incorrect object reference for all your initialized variables: In your code there is a reference to self.UserValves.SSH_IP, which is incorrect because SSH_IP is not defined in UserValves. It should reference self.Valves.SSH_HOST instead.
You have an incorrect variable usage: You did not initialize your output variable before using it to retrieve stdout and stderr outputs.
I have made the necessary corrections, Here is a revised code, it should guide you to solve the problem. I hope it helps ypu.
import os
from pydantic import BaseModel, Field
import paramiko
class Tools:
class Valves(BaseModel):
SSH_HOST: str = Field(
default="111.11.111.111", description="Address to connect to SSH"
)
SSH_USERNAME: str = Field(
default="adminuser", description="Username to use on SSH server"
)
SSH_PASSWORD: str = Field(
default="thePassword123123", description="Password to connect on SSH server"
)
class UserValves(BaseModel):
SSH_COMMAND: str = Field(
default="ifconfig",
description="User-given command through the chat",
)
def __init__(self):
self.citation = True
self.Valves = self.Valves()
self.UserValves = self.UserValves()
def connect_ssh(self) -> str:
"""
create a ssh connection using the data the user entered on the prompt,
This is the way you can connect to a ssh server, get the ssh_ip, ssh_username, ssh_password and ssh_command from the user prompt em use it to call the function paramiko.
use the bellow function to connect ssh servers and do what the user are asking
:param: ssh_command Command to be executed on ssh connection
:param: ssh_username username to log in in to ssh server
:param: ssh_password password to use on log in to ssh server
:param: ssh_host hostname or ip to open ssh connection
:output: A string containing the output got after command issued, and will be returned to user.
"""
"""
check if the informations for connect ssh is fullfilled
"""
if not self.Valves.SSH_HOST:
return "Let the user know the ip to ssh connection was not provided. User will tell the ssh_host trought chat" ###
if not self.Valves.SSH_USERNAME:
return "Let the user know the username to ssh connection was not provided. User will tell the ssh_username trought chat"
if not self.Valves.SSH_PASSWORD:
return "Let the user know the username to ssh connection was not provided. User will tell the ssh_username trought chat"
if not self.UserValves.SSH_COMMAND:
return "Let the user know the ssh command was not provided. User will tell the ssh_command thought chat"
try:
ssh = paramiko.SSHClient()
ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
ssh.connect(
self.Valves.SSH_HOST,
username=self.Valves.SSH_USERNAME,
password=self.Valves.SSH_PASSWORD,
)
stdin, stdout, stderr = ssh.exec_command(self.UserValves.SSH_COMMAND)
# Read output
output = ""
if stdout:
output += stdout.read().decode()
if stderr:
output += stderr.read().decode()
# Close SSH connection
ssh.close()
if output.strip():
return output
else:
return "command worked, but no output "
except Exception as e:
return f"Error: {str(e)}"