fluzzi32
12543c683e
pointing to /config. This avoids running the conn command during the build process and ensures a
cleaner setup.
2. Copilot UI Fix: Resolved a double-escaping bug in the terminal bottom bar. Device prompts (like
6WIND-PE1>) will now render correctly instead of showing HTML entities like >.
3. AI Model Update: Updated the default engineer model in connpy/ai.py to
gemini/gemini-3.1-flash-lite, removing the deprecated -preview suffix.
4. Standardized Timeouts: Unified all default timeouts to 20 seconds across the board. This includes
direct execution (run/test), modern playbooks (v2), and classic task-based playbooks (v1).
5. Documentation Update: Regenerated the full documentation site in the docs/ directory using pdoc to
reflect the latest changes.
6. Cleanup: Removed all debug prints from connpy/core.py and restored the docker/logs/.gitignore
file.
6214 lines
315 KiB
HTML
6214 lines
315 KiB
HTML
<!doctype html>
|
||
<html lang="en">
|
||
<head>
|
||
<meta charset="utf-8">
|
||
<meta name="viewport" content="width=device-width, initial-scale=1, minimum-scale=1">
|
||
<meta name="generator" content="pdoc3 0.11.5">
|
||
<title>connpy API documentation</title>
|
||
<meta name="description" content="<p align="center">
|
||
<img src="https://nginx.gederico.dynu.net/images/CONNPY-resized.png" alt="App Logo">
|
||
</p> …">
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/sanitize.min.css" integrity="sha512-y1dtMcuvtTMJc1yPgEqF0ZjQbhnc/bFhyvIyVNb9Zk5mIGtqVaAB1Ttl28su8AvFMOY0EwRbAe+HCLqj6W7/KA==" crossorigin>
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/10up-sanitize.css/13.0.0/typography.min.css" integrity="sha512-Y1DYSb995BAfxobCkKepB1BqJJTPrOp3zPL74AWFugHHmmdcvO+C48WLrUOlhGMc0QG7AE3f7gmvvcrmX2fDoA==" crossorigin>
|
||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/default.min.css" crossorigin>
|
||
<style>:root{--highlight-color:#fe9}.flex{display:flex !important}body{line-height:1.5em}#content{padding:20px}#sidebar{padding:1.5em;overflow:hidden}#sidebar > *:last-child{margin-bottom:2cm}.http-server-breadcrumbs{font-size:130%;margin:0 0 15px 0}#footer{font-size:.75em;padding:5px 30px;border-top:1px solid #ddd;text-align:right}#footer p{margin:0 0 0 1em;display:inline-block}#footer p:last-child{margin-right:30px}h1,h2,h3,h4,h5{font-weight:300}h1{font-size:2.5em;line-height:1.1em}h2{font-size:1.75em;margin:2em 0 .50em 0}h3{font-size:1.4em;margin:1.6em 0 .7em 0}h4{margin:0;font-size:105%}h1:target,h2:target,h3:target,h4:target,h5:target,h6:target{background:var(--highlight-color);padding:.2em 0}a{color:#058;text-decoration:none;transition:color .2s ease-in-out}a:visited{color:#503}a:hover{color:#b62}.title code{font-weight:bold}h2[id^="header-"]{margin-top:2em}.ident{color:#900;font-weight:bold}pre code{font-size:.8em;line-height:1.4em;padding:1em;display:block}code{background:#f3f3f3;font-family:"DejaVu Sans Mono",monospace;padding:1px 4px;overflow-wrap:break-word}h1 code{background:transparent}pre{border-top:1px solid #ccc;border-bottom:1px solid #ccc;margin:1em 0}#http-server-module-list{display:flex;flex-flow:column}#http-server-module-list div{display:flex}#http-server-module-list dt{min-width:10%}#http-server-module-list p{margin-top:0}.toc ul,#index{list-style-type:none;margin:0;padding:0}#index code{background:transparent}#index h3{border-bottom:1px solid #ddd}#index ul{padding:0}#index h4{margin-top:.6em;font-weight:bold}@media (min-width:200ex){#index .two-column{column-count:2}}@media (min-width:300ex){#index .two-column{column-count:3}}dl{margin-bottom:2em}dl dl:last-child{margin-bottom:4em}dd{margin:0 0 1em 3em}#header-classes + dl > dd{margin-bottom:3em}dd dd{margin-left:2em}dd p{margin:10px 0}.name{background:#eee;font-size:.85em;padding:5px 10px;display:inline-block;min-width:40%}.name:hover{background:#e0e0e0}dt:target .name{background:var(--highlight-color)}.name > span:first-child{white-space:nowrap}.name.class > span:nth-child(2){margin-left:.4em}.inherited{color:#999;border-left:5px solid #eee;padding-left:1em}.inheritance em{font-style:normal;font-weight:bold}.desc h2{font-weight:400;font-size:1.25em}.desc h3{font-size:1em}.desc dt code{background:inherit}.source > summary,.git-link-div{color:#666;text-align:right;font-weight:400;font-size:.8em;text-transform:uppercase}.source summary > *{white-space:nowrap;cursor:pointer}.git-link{color:inherit;margin-left:1em}.source pre{max-height:500px;overflow:auto;margin:0}.source pre code{font-size:12px;overflow:visible;min-width:max-content}.hlist{list-style:none}.hlist li{display:inline}.hlist li:after{content:',\2002'}.hlist li:last-child:after{content:none}.hlist .hlist{display:inline;padding-left:1em}img{max-width:100%}td{padding:0 .5em}.admonition{padding:.1em 1em;margin:1em 0}.admonition-title{font-weight:bold}.admonition.note,.admonition.info,.admonition.important{background:#aef}.admonition.todo,.admonition.versionadded,.admonition.tip,.admonition.hint{background:#dfd}.admonition.warning,.admonition.versionchanged,.admonition.deprecated{background:#fd4}.admonition.error,.admonition.danger,.admonition.caution{background:lightpink}</style>
|
||
<style media="screen and (min-width: 700px)">@media screen and (min-width:700px){#sidebar{width:30%;height:100vh;overflow:auto;position:sticky;top:0}#content{width:70%;max-width:100ch;padding:3em 4em;border-left:1px solid #ddd}pre code{font-size:1em}.name{font-size:1em}main{display:flex;flex-direction:row-reverse;justify-content:flex-end}.toc ul ul,#index ul ul{padding-left:1em}.toc > ul > li{margin-top:.5em}}</style>
|
||
<style media="print">@media print{#sidebar h1{page-break-before:always}.source{display:none}}@media print{*{background:transparent !important;color:#000 !important;box-shadow:none !important;text-shadow:none !important}a[href]:after{content:" (" attr(href) ")";font-size:90%}a[href][title]:after{content:none}abbr[title]:after{content:" (" attr(title) ")"}.ir a:after,a[href^="javascript:"]:after,a[href^="#"]:after{content:""}pre,blockquote{border:1px solid #999;page-break-inside:avoid}thead{display:table-header-group}tr,img{page-break-inside:avoid}img{max-width:100% !important}@page{margin:0.5cm}p,h2,h3{orphans:3;widows:3}h1,h2,h3,h4,h5,h6{page-break-after:avoid}}</style>
|
||
<script defer src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js" integrity="sha512-D9gUyxqja7hBtkWpPWGt9wfbfaMGVt9gnyCvYa+jojwwPHLCzUm5i8rpk7vD7wNee9bA35eYIjobYPaQuKS1MQ==" crossorigin></script>
|
||
<script>window.addEventListener('DOMContentLoaded', () => {
|
||
hljs.configure({languages: ['bash', 'css', 'diff', 'graphql', 'ini', 'javascript', 'json', 'plaintext', 'python', 'python-repl', 'rust', 'shell', 'sql', 'typescript', 'xml', 'yaml']});
|
||
hljs.highlightAll();
|
||
/* Collapse source docstrings */
|
||
setTimeout(() => {
|
||
[...document.querySelectorAll('.hljs.language-python > .hljs-string')]
|
||
.filter(el => el.innerHTML.length > 200 && ['"""', "'''"].includes(el.innerHTML.substring(0, 3)))
|
||
.forEach(el => {
|
||
let d = document.createElement('details');
|
||
d.classList.add('hljs-string');
|
||
d.innerHTML = '<summary>"""</summary>' + el.innerHTML.substring(3);
|
||
el.replaceWith(d);
|
||
});
|
||
}, 100);
|
||
})</script>
|
||
</head>
|
||
<body>
|
||
<main>
|
||
<article id="content">
|
||
<header>
|
||
<h1 class="title">Package <code>connpy</code></h1>
|
||
</header>
|
||
<section id="section-intro">
|
||
<p align="center">
|
||
<img src="https://nginx.gederico.dynu.net/images/CONNPY-resized.png" alt="App Logo">
|
||
</p>
|
||
<h1 id="connpy">Connpy</h1>
|
||
<p><a href="https://pypi.org/pypi/connpy/"><img alt="" src="https://img.shields.io/pypi/v/connpy.svg?style=flat-square"></a>
|
||
<a href="https://pypi.org/pypi/connpy/"><img alt="" src="https://img.shields.io/pypi/pyversions/connpy.svg?style=flat-square"></a>
|
||
<a href="https://github.com/fluzzi/connpy/blob/main/LICENSE"><img alt="" src="https://img.shields.io/pypi/l/connpy.svg?style=flat-square"></a>
|
||
<a href="https://pypi.org/pypi/connpy/"><img alt="" src="https://img.shields.io/pypi/dm/connpy.svg?style=flat-square"></a></p>
|
||
<p><strong>Connpy</strong> is a powerful Connection Manager and Network Automation Platform for Linux, Mac, and Docker. It provides a unified interface for <strong>SSH, SFTP, Telnet, kubectl, Docker pods, and AWS SSM</strong>.</p>
|
||
<p>The v6 release introduces the <strong>AI Copilot</strong>, an interactive terminal assistant that understands your network context and helps you manage your infrastructure more intelligently.</p>
|
||
<h2 id="ai-copilot-new-in-v6">🤖 AI Copilot (New in v6)</h2>
|
||
<p>The AI Copilot is deeply integrated into your terminal workflow:
|
||
- <strong>Terminal Context Awareness</strong>: The Copilot can "see" your screen output, helping you diagnose errors or analyze command results in real-time.
|
||
- <strong>Hybrid Multi-Agent System</strong>: Automatically escalates complex tasks between the <strong>Network Engineer</strong> (execution) and the <strong>Network Architect</strong> (strategy).
|
||
- <strong>MCP Integration</strong>: Dynamically load tools from external providers (6WIND, AWS, etc.) via the Model Context Protocol.
|
||
- <strong>Interactive Chat</strong>: Launch with <code>conn <a title="connpy.ai" href="#connpy.ai">ai</a></code> for a collaborative troubleshooting session.</p>
|
||
<h2 id="core-features">Core Features</h2>
|
||
<ul>
|
||
<li><strong>Multi-Protocol</strong>: Native support for SSH, SFTP, Telnet, kubectl, Docker exec, and AWS SSM.</li>
|
||
<li><strong>Context Management</strong>: Set regex-based contexts to manage specific nodes across different environments (work, home, clients).</li>
|
||
<li><strong>Advanced Inventory</strong>:<ul>
|
||
<li>Organize nodes in folders (<code>@folder</code>) and subfolders (<code>@subfolder@folder</code>).</li>
|
||
<li>Use Global Profiles (<code>@profilename</code>) to manage shared credentials easily.</li>
|
||
<li>Bulk creation, copying, moving, and export/import of nodes.</li>
|
||
</ul>
|
||
</li>
|
||
<li><strong>Modern UI</strong>: High-performance terminal experience with <code>prompt-toolkit</code>, including:<ul>
|
||
<li>Fuzzy search integration with <code>fzf</code>.</li>
|
||
<li>Advanced tab completion.</li>
|
||
<li>Syntax highlighting and customizable themes.</li>
|
||
</ul>
|
||
</li>
|
||
<li><strong>Automation Engine</strong>: Run parallel tasks and playbooks on multiple devices with variable support.</li>
|
||
<li><strong>Plugin System</strong>: Build and execute custom Python scripts locally or on a remote gRPC server.</li>
|
||
<li><strong>gRPC Architecture</strong>: Fully decoupled Client/Server model for distributed management.</li>
|
||
<li><strong>Privacy & Sync</strong>: Local-first encrypted storage (RSA/OAEP) with optional Google Drive backup.</li>
|
||
</ul>
|
||
<h2 id="installation">Installation</h2>
|
||
<pre><code class="language-bash">pip install connpy
|
||
</code></pre>
|
||
<h3 id="run-it-in-windowslinux-using-docker">Run it in Windows/Linux using Docker</h3>
|
||
<pre><code class="language-bash">git clone https://github.com/fluzzi/connpy
|
||
cd connpy
|
||
docker compose build
|
||
|
||
# Run it like a native app (completely silent)
|
||
docker compose --log-level ERROR run --rm --remove-orphans connpy-app [command]
|
||
|
||
# Pro Tip: Add this alias for a 100% native experience from any folder
|
||
alias conn='docker compose -f /path/to/connpy/docker-compose.yml --log-level ERROR run --rm --remove-orphans connpy-app'
|
||
</code></pre>
|
||
<hr>
|
||
<h2 id="privacy-integration">🔒 Privacy & Integration</h2>
|
||
<h3 id="privacy-policy">Privacy Policy</h3>
|
||
<p>Connpy is committed to protecting your privacy:
|
||
- <strong>Local Storage</strong>: All server addresses, usernames, and passwords are encrypted and stored <strong>only</strong> on your machine. No data is transmitted to our servers.
|
||
- <strong>Data Access</strong>: Data is used solely for managing and automating your connections.</p>
|
||
<h3 id="google-integration">Google Integration</h3>
|
||
<p>Used strictly for backup:
|
||
- <strong>Backup</strong>: Sync your encrypted configuration with your Google Drive account.
|
||
- <strong>Scoped Access</strong>: Connpy only accesses its own backup files.</p>
|
||
<hr>
|
||
<h2 id="usage">Usage</h2>
|
||
<pre><code class="language-bash">usage: conn [-h] [--add | --del | --mod | --show | --debug] [node|folder] [--sftp]
|
||
conn {profile,move,copy,list,bulk,export,import,ai,run,api,plugin,config,sync,context} ...
|
||
</code></pre>
|
||
<h3 id="basic-examples">Basic Examples:</h3>
|
||
<pre><code class="language-bash"># Add a folder and subfolder
|
||
conn --add @office
|
||
conn --add @datacenter@office
|
||
|
||
# Add a node with a profile
|
||
conn --add server1@datacenter@office --profile @myuser
|
||
|
||
# Connect to a node (fuzzy match)
|
||
conn server1
|
||
|
||
# Start the AI Copilot
|
||
conn ai
|
||
|
||
# Run a command on all nodes in a folder
|
||
conn run @office "uptime"
|
||
</code></pre>
|
||
<hr>
|
||
<h2 id="plugin-requirements-for-connpy">Plugin Requirements for Connpy</h2>
|
||
<h3 id="remote-plugin-execution">Remote Plugin Execution</h3>
|
||
<p>When Connpy operates in remote mode, plugins are executed <strong>transparently on the server</strong>:
|
||
- The client automatically downloads the plugin source code (<code>Parser</code> class context) to generate the local <code>argparse</code> structure and provide autocompletion.
|
||
- The execution phase (<code>Entrypoint</code> class) is redirected via gRPC streams to execute in the server's memory.
|
||
- You can manage remote plugins using the <code>--remote</code> flag.</p>
|
||
<h3 id="general-structure">General Structure</h3>
|
||
<ul>
|
||
<li>The plugin script must define specific classes:</li>
|
||
<li><strong>Class <code>Parser</code></strong>: Handles <code>argparse.ArgumentParser</code> initialization.</li>
|
||
<li><strong>Class <code>Entrypoint</code></strong>: Main execution logic (receives <code>args</code>, <code>parser</code>, and <code>connapp</code>).</li>
|
||
<li><strong>Class <code>Preload</code></strong>: (Optional) For modifying core app behavior or registering hooks.</li>
|
||
</ul>
|
||
<h3 id="preload-modifications-and-hooks">Preload Modifications and Hooks</h3>
|
||
<p>You can customize the behavior of core classes using hooks:
|
||
- <strong><code>modify(method)</code></strong>: Alter class instances (e.g., <code>connapp.config</code>, <code>connapp.ai</code>).
|
||
- <strong><code>register_pre_hook(method)</code></strong>: Logic to run before a method execution.
|
||
- <strong><code>register_post_hook(method)</code></strong>: Logic to run after a method execution.</p>
|
||
<h3 id="command-completion-support">Command Completion Support</h3>
|
||
<p>Plugins can provide intelligent tab completion:
|
||
1. <strong>Tree-based Completion (Recommended)</strong>: Define <code>_connpy_tree(info)</code> returning a navigation dictionary.
|
||
2. <strong>Legacy Completion</strong>: Define <code>_connpy_completion(wordsnumber, words, info)</code>.</p>
|
||
<hr>
|
||
<h2 id="grpc-service-architecture">⚙️ gRPC Service Architecture</h2>
|
||
<p>Connpy can operate in a decoupled mode:
|
||
1. <strong>Start the API (Server)</strong>: <code>conn api -s 50051</code>
|
||
2. <strong>Configure the Client</strong>:
|
||
<code>bash
|
||
conn config --service-mode remote
|
||
conn config --remote-host localhost:50051</code>
|
||
All inventory management and execution will now happen on the server.</p>
|
||
<hr>
|
||
<h2 id="automation-module-api">🐍 Automation Module (API)</h2>
|
||
<p>You can use <code><a title="connpy" href="#connpy">connpy</a></code> as a Python library for your own scripts.</p>
|
||
<h3 id="basic-execution">Basic Execution</h3>
|
||
<pre><code class="language-python">import connpy
|
||
router = connpy.node("uniqueName", "1.1.1.1", user="admin")
|
||
router.run(["show ip int brief"])
|
||
print(router.output)
|
||
</code></pre>
|
||
<h3 id="parallel-tasks-with-variables">Parallel Tasks with Variables</h3>
|
||
<pre><code class="language-python">import connpy
|
||
config = connpy.configfile()
|
||
nodes = config.getitem("@office", ["router1", "router2"])
|
||
routers = connpy.nodes(nodes, config=config)
|
||
|
||
variables = {
|
||
"router1@office": {"id": "1"},
|
||
"__global__": {"mask": "255.255.255.0"}
|
||
}
|
||
routers.run(["interface lo{id}", "ip address 10.0.0.{id} {mask}"], variables)
|
||
</code></pre>
|
||
<h3 id="ai-programmatic-use">AI Programmatic Use</h3>
|
||
<pre><code class="language-python">import connpy
|
||
myai = connpy.ai(connpy.configfile())
|
||
response = myai.ask("What is the status of the BGP neighbors in the office?")
|
||
</code></pre>
|
||
<hr>
|
||
<p><em>For detailed developer notes and plugin hooks documentation, see the <a href="https://fluzzi.github.io/connpy/">Documentation</a>.</em></p>
|
||
</section>
|
||
<section>
|
||
<h2 class="section-title" id="header-submodules">Sub-modules</h2>
|
||
<dl>
|
||
<dt><code class="name"><a title="connpy.cli" href="cli/index.html">connpy.cli</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.grpc_layer" href="grpc_layer/index.html">connpy.grpc_layer</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.mcp_client" href="mcp_client.html">connpy.mcp_client</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.proto" href="proto/index.html">connpy.proto</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.services" href="services/index.html">connpy.services</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.tests" href="tests/index.html">connpy.tests</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.tunnels" href="tunnels.html">connpy.tunnels</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt><code class="name"><a title="connpy.utils" href="utils.html">connpy.utils</a></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
</dl>
|
||
</section>
|
||
<section>
|
||
</section>
|
||
<section>
|
||
</section>
|
||
<section>
|
||
<h2 class="section-title" id="header-classes">Classes</h2>
|
||
<dl>
|
||
<dt id="connpy.Plugins"><code class="flex name class">
|
||
<span>class <span class="ident">Plugins</span></span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">class Plugins:
|
||
def __init__(self):
|
||
self.plugins = {}
|
||
self.plugin_parsers = {}
|
||
self.preloads = {}
|
||
self.remote_plugins = {}
|
||
self.preferences = {}
|
||
|
||
def _load_preferences(self, config_dir):
|
||
import json
|
||
path = os.path.join(config_dir, "plugin_preferences.json")
|
||
try:
|
||
with open(path) as f:
|
||
self.preferences = json.load(f)
|
||
except (FileNotFoundError, json.JSONDecodeError):
|
||
self.preferences = {}
|
||
|
||
def _save_preferences(self, config_dir):
|
||
import json
|
||
path = os.path.join(config_dir, "plugin_preferences.json")
|
||
try:
|
||
with open(path, "w") as f:
|
||
json.dump(self.preferences, f, indent=4)
|
||
except OSError as e:
|
||
printer.error(f"Failed to save plugin preferences: {e}")
|
||
|
||
|
||
def verify_script(self, file_path):
|
||
"""
|
||
Verifies that a given Python script meets specific structural requirements.
|
||
|
||
This function checks a Python script for compliance with predefined structural
|
||
rules. It ensures that the script contains only allowed top-level elements
|
||
(functions, classes, imports, pass statements, and a specific if __name__ block)
|
||
and that it includes mandatory classes with specific attributes and methods.
|
||
|
||
### Arguments:
|
||
- file_path (str): The file path of the Python script to be verified.
|
||
|
||
### Returns:
|
||
- str: A message indicating the type of violation if the script doesn't meet
|
||
the requirements, or False if all requirements are met.
|
||
|
||
### Verifications:
|
||
- The presence of only allowed top-level elements.
|
||
- The existence of two specific classes: 'Parser' and 'Entrypoint'. and/or specific class: Preload.
|
||
- 'Parser' class must only have an '__init__' method and must assign 'self.parser'.
|
||
- 'Entrypoint' class must have an '__init__' method accepting specific arguments.
|
||
|
||
If any of these checks fail, the function returns an error message indicating
|
||
the reason. If the script passes all checks, the function returns False,
|
||
indicating successful verification.
|
||
|
||
### Exceptions:
|
||
- SyntaxError: If the script contains a syntax error, it is caught and
|
||
returned as a part of the error message.
|
||
"""
|
||
with open(file_path, 'r') as file:
|
||
source_code = file.read()
|
||
|
||
try:
|
||
tree = ast.parse(source_code)
|
||
except SyntaxError as e:
|
||
return f"Syntax error in file: {e}"
|
||
|
||
|
||
has_parser = False
|
||
has_entrypoint = False
|
||
has_preload = False
|
||
|
||
for node in tree.body:
|
||
# Allow only function definitions, class definitions, and pass statements at top-level
|
||
if isinstance(node, ast.If):
|
||
# Check for the 'if __name__ == "__main__":' block
|
||
if not (isinstance(node.test, ast.Compare) and
|
||
isinstance(node.test.left, ast.Name) and
|
||
node.test.left.id == '__name__' and
|
||
((hasattr(ast, 'Str') and isinstance(node.test.comparators[0], getattr(ast, 'Str')) and node.test.comparators[0].s == '__main__') or
|
||
(hasattr(ast, 'Constant') and isinstance(node.test.comparators[0], getattr(ast, 'Constant')) and node.test.comparators[0].value == '__main__'))):
|
||
return "Only __name__ == __main__ If is allowed"
|
||
|
||
elif not isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.Import, ast.ImportFrom, ast.Pass)):
|
||
return f"Plugin can only have pass, functions, classes and imports. {node} is not allowed" # Reject any other AST types
|
||
|
||
if isinstance(node, ast.ClassDef):
|
||
|
||
if node.name == 'Parser':
|
||
has_parser = True
|
||
# Ensure Parser class has only the __init__ method and assigns self.parser
|
||
if not all(isinstance(method, ast.FunctionDef) and method.name == '__init__' for method in node.body):
|
||
return "Parser class should only have __init__ method"
|
||
|
||
# Check if 'self.parser' is assigned in __init__ method
|
||
init_method = node.body[0]
|
||
assigned_attrs = [target.attr for expr in init_method.body if isinstance(expr, ast.Assign) for target in expr.targets if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == 'self']
|
||
if 'parser' not in assigned_attrs:
|
||
return "Parser class should set self.parser"
|
||
|
||
|
||
elif node.name == 'Entrypoint':
|
||
has_entrypoint = True
|
||
init_method = next((item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == '__init__'), None)
|
||
if not init_method or len(init_method.args.args) != 4: # self, args, parser, conapp
|
||
return "Entrypoint class should have method __init__ and accept only arguments: args, parser and connapp" # 'Entrypoint' __init__ does not have correct signature
|
||
|
||
elif node.name == 'Preload':
|
||
has_preload = True
|
||
init_method = next((item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == '__init__'), None)
|
||
if not init_method or len(init_method.args.args) != 2: # self, connapp
|
||
return "Preload class should have method __init__ and accept only argument: connapp" # 'Preload' __init__ does not have correct signature
|
||
|
||
# Applying the combination logic based on class presence
|
||
if has_parser and not has_entrypoint:
|
||
return "Parser requires Entrypoint class to be present."
|
||
elif has_entrypoint and not has_parser:
|
||
return "Entrypoint requires Parser class to be present."
|
||
|
||
if not (has_parser or has_entrypoint or has_preload):
|
||
return "No valid class (Parser, Entrypoint, or Preload) found."
|
||
|
||
return False # All requirements met, no error
|
||
|
||
def _import_from_path(self, path):
|
||
spec = importlib.util.spec_from_file_location("module.name", path)
|
||
module = importlib.util.module_from_spec(spec)
|
||
sys.modules["module.name"] = module
|
||
spec.loader.exec_module(module)
|
||
return module
|
||
|
||
def _import_plugins_to_argparse(self, directory, subparsers, remote_enabled=False):
|
||
if not os.path.exists(directory):
|
||
return
|
||
for filename in os.listdir(directory):
|
||
commands = subparsers.choices.keys()
|
||
if filename.endswith(".py"):
|
||
root_filename = os.path.splitext(filename)[0]
|
||
if root_filename in commands:
|
||
continue
|
||
|
||
# Check preferences: if remote is preferred AND remote is enabled, skip local loading
|
||
if remote_enabled and self.preferences.get(root_filename) == "remote":
|
||
continue
|
||
|
||
# Construct the full path
|
||
filepath = os.path.join(directory, filename)
|
||
check_file = self.verify_script(filepath)
|
||
if check_file:
|
||
printer.error(f"Failed to load plugin: {filename}. Reason: {check_file}")
|
||
continue
|
||
else:
|
||
self.plugins[root_filename] = self._import_from_path(filepath)
|
||
if hasattr(self.plugins[root_filename], "Parser"):
|
||
self.plugin_parsers[root_filename] = self.plugins[root_filename].Parser()
|
||
plugin = self.plugin_parsers[root_filename]
|
||
# Default to RichHelpFormatter if plugin doesn't set one
|
||
try:
|
||
from rich_argparse import RichHelpFormatter as _RHF
|
||
fmt = plugin.parser.formatter_class
|
||
if fmt is argparse.HelpFormatter or fmt is argparse.RawTextHelpFormatter or fmt is argparse.RawDescriptionHelpFormatter:
|
||
fmt = _RHF
|
||
except ImportError:
|
||
fmt = plugin.parser.formatter_class
|
||
subparsers.add_parser(root_filename, parents=[self.plugin_parsers[root_filename].parser], add_help=False, help=plugin.parser.description, usage=plugin.parser.usage, description=plugin.parser.description, epilog=plugin.parser.epilog, formatter_class=fmt)
|
||
if hasattr(self.plugins[root_filename], "Preload"):
|
||
self.preloads[root_filename] = self.plugins[root_filename]
|
||
|
||
def _import_remote_plugins_to_argparse(self, plugin_stub, subparsers, cache_dir, force_sync=False):
|
||
import hashlib
|
||
os.makedirs(cache_dir, exist_ok=True)
|
||
|
||
try:
|
||
remote_plugins_info = plugin_stub.list_plugins()
|
||
except Exception:
|
||
return
|
||
|
||
# Pruning: Remove local cached files that are no longer on the server
|
||
for local_file in os.listdir(cache_dir):
|
||
if local_file.endswith(".py"):
|
||
name = local_file[:-3]
|
||
if name not in remote_plugins_info:
|
||
try:
|
||
os.remove(os.path.join(cache_dir, local_file))
|
||
except Exception:
|
||
pass
|
||
|
||
for name, info in remote_plugins_info.items():
|
||
if not info.get("enabled", True):
|
||
continue
|
||
|
||
pref = self.preferences.get(name, "local")
|
||
if pref != "remote" and name in self.plugins:
|
||
continue
|
||
if not force_sync and name in subparsers.choices:
|
||
continue
|
||
|
||
cache_path = os.path.join(cache_dir, f"{name}.py")
|
||
|
||
# Hash comparison
|
||
remote_hash = info.get("hash", "")
|
||
local_hash = ""
|
||
if os.path.exists(cache_path):
|
||
try:
|
||
with open(cache_path, "rb") as f:
|
||
local_hash = hashlib.md5(f.read()).hexdigest()
|
||
except Exception:
|
||
pass
|
||
|
||
# Update only if hash differs or force_sync is True
|
||
if force_sync or remote_hash != local_hash or not os.path.exists(cache_path):
|
||
try:
|
||
source = plugin_stub.get_plugin_source(name)
|
||
with open(cache_path, "w") as f:
|
||
f.write(source)
|
||
except Exception as e:
|
||
printer.warning(f"Failed to sync remote plugin {name}: {e}")
|
||
continue
|
||
|
||
# Verify and load
|
||
check_file = self.verify_script(cache_path)
|
||
if check_file:
|
||
printer.warning(f"Remote plugin {name} failed verification: {check_file}")
|
||
continue
|
||
|
||
module = self._import_from_path(cache_path)
|
||
if hasattr(module, "Parser"):
|
||
self.plugin_parsers[name] = module.Parser()
|
||
self.remote_plugins[name] = True
|
||
plugin = self.plugin_parsers[name]
|
||
try:
|
||
from rich_argparse import RichHelpFormatter as _RHF
|
||
fmt = plugin.parser.formatter_class
|
||
if fmt is argparse.HelpFormatter or fmt is argparse.RawTextHelpFormatter or fmt is argparse.RawDescriptionHelpFormatter:
|
||
fmt = _RHF
|
||
except ImportError:
|
||
fmt = plugin.parser.formatter_class
|
||
|
||
# If force_sync, we might be re-registering, but argparse subparsers.add_parser
|
||
# might fail if it exists. We check if it's already there.
|
||
if name not in subparsers.choices:
|
||
subparsers.add_parser(
|
||
name,
|
||
parents=[plugin.parser],
|
||
add_help=False,
|
||
help=f"[remote] {plugin.parser.description}",
|
||
usage=plugin.parser.usage,
|
||
description=plugin.parser.description,
|
||
epilog=plugin.parser.epilog,
|
||
formatter_class=fmt
|
||
)</code></pre>
|
||
</details>
|
||
<div class="desc"></div>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="connpy.Plugins.verify_script"><code class="name flex">
|
||
<span>def <span class="ident">verify_script</span></span>(<span>self, file_path)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def verify_script(self, file_path):
|
||
"""
|
||
Verifies that a given Python script meets specific structural requirements.
|
||
|
||
This function checks a Python script for compliance with predefined structural
|
||
rules. It ensures that the script contains only allowed top-level elements
|
||
(functions, classes, imports, pass statements, and a specific if __name__ block)
|
||
and that it includes mandatory classes with specific attributes and methods.
|
||
|
||
### Arguments:
|
||
- file_path (str): The file path of the Python script to be verified.
|
||
|
||
### Returns:
|
||
- str: A message indicating the type of violation if the script doesn't meet
|
||
the requirements, or False if all requirements are met.
|
||
|
||
### Verifications:
|
||
- The presence of only allowed top-level elements.
|
||
- The existence of two specific classes: 'Parser' and 'Entrypoint'. and/or specific class: Preload.
|
||
- 'Parser' class must only have an '__init__' method and must assign 'self.parser'.
|
||
- 'Entrypoint' class must have an '__init__' method accepting specific arguments.
|
||
|
||
If any of these checks fail, the function returns an error message indicating
|
||
the reason. If the script passes all checks, the function returns False,
|
||
indicating successful verification.
|
||
|
||
### Exceptions:
|
||
- SyntaxError: If the script contains a syntax error, it is caught and
|
||
returned as a part of the error message.
|
||
"""
|
||
with open(file_path, 'r') as file:
|
||
source_code = file.read()
|
||
|
||
try:
|
||
tree = ast.parse(source_code)
|
||
except SyntaxError as e:
|
||
return f"Syntax error in file: {e}"
|
||
|
||
|
||
has_parser = False
|
||
has_entrypoint = False
|
||
has_preload = False
|
||
|
||
for node in tree.body:
|
||
# Allow only function definitions, class definitions, and pass statements at top-level
|
||
if isinstance(node, ast.If):
|
||
# Check for the 'if __name__ == "__main__":' block
|
||
if not (isinstance(node.test, ast.Compare) and
|
||
isinstance(node.test.left, ast.Name) and
|
||
node.test.left.id == '__name__' and
|
||
((hasattr(ast, 'Str') and isinstance(node.test.comparators[0], getattr(ast, 'Str')) and node.test.comparators[0].s == '__main__') or
|
||
(hasattr(ast, 'Constant') and isinstance(node.test.comparators[0], getattr(ast, 'Constant')) and node.test.comparators[0].value == '__main__'))):
|
||
return "Only __name__ == __main__ If is allowed"
|
||
|
||
elif not isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.Import, ast.ImportFrom, ast.Pass)):
|
||
return f"Plugin can only have pass, functions, classes and imports. {node} is not allowed" # Reject any other AST types
|
||
|
||
if isinstance(node, ast.ClassDef):
|
||
|
||
if node.name == 'Parser':
|
||
has_parser = True
|
||
# Ensure Parser class has only the __init__ method and assigns self.parser
|
||
if not all(isinstance(method, ast.FunctionDef) and method.name == '__init__' for method in node.body):
|
||
return "Parser class should only have __init__ method"
|
||
|
||
# Check if 'self.parser' is assigned in __init__ method
|
||
init_method = node.body[0]
|
||
assigned_attrs = [target.attr for expr in init_method.body if isinstance(expr, ast.Assign) for target in expr.targets if isinstance(target, ast.Attribute) and isinstance(target.value, ast.Name) and target.value.id == 'self']
|
||
if 'parser' not in assigned_attrs:
|
||
return "Parser class should set self.parser"
|
||
|
||
|
||
elif node.name == 'Entrypoint':
|
||
has_entrypoint = True
|
||
init_method = next((item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == '__init__'), None)
|
||
if not init_method or len(init_method.args.args) != 4: # self, args, parser, conapp
|
||
return "Entrypoint class should have method __init__ and accept only arguments: args, parser and connapp" # 'Entrypoint' __init__ does not have correct signature
|
||
|
||
elif node.name == 'Preload':
|
||
has_preload = True
|
||
init_method = next((item for item in node.body if isinstance(item, ast.FunctionDef) and item.name == '__init__'), None)
|
||
if not init_method or len(init_method.args.args) != 2: # self, connapp
|
||
return "Preload class should have method __init__ and accept only argument: connapp" # 'Preload' __init__ does not have correct signature
|
||
|
||
# Applying the combination logic based on class presence
|
||
if has_parser and not has_entrypoint:
|
||
return "Parser requires Entrypoint class to be present."
|
||
elif has_entrypoint and not has_parser:
|
||
return "Entrypoint requires Parser class to be present."
|
||
|
||
if not (has_parser or has_entrypoint or has_preload):
|
||
return "No valid class (Parser, Entrypoint, or Preload) found."
|
||
|
||
return False # All requirements met, no error</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Verifies that a given Python script meets specific structural requirements.</p>
|
||
<p>This function checks a Python script for compliance with predefined structural
|
||
rules. It ensures that the script contains only allowed top-level elements
|
||
(functions, classes, imports, pass statements, and a specific if <strong>name</strong> block)
|
||
and that it includes mandatory classes with specific attributes and methods.</p>
|
||
<h3 id="arguments">Arguments:</h3>
|
||
<pre><code>- file_path (str): The file path of the Python script to be verified.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>- str: A message indicating the type of violation if the script doesn't meet
|
||
the requirements, or False if all requirements are met.
|
||
</code></pre>
|
||
<h3 id="verifications">Verifications:</h3>
|
||
<pre><code>- The presence of only allowed top-level elements.
|
||
- The existence of two specific classes: 'Parser' and 'Entrypoint'. and/or specific class: Preload.
|
||
- 'Parser' class must only have an '__init__' method and must assign 'self.parser'.
|
||
- 'Entrypoint' class must have an '__init__' method accepting specific arguments.
|
||
</code></pre>
|
||
<p>If any of these checks fail, the function returns an error message indicating
|
||
the reason. If the script passes all checks, the function returns False,
|
||
indicating successful verification.</p>
|
||
<h3 id="exceptions">Exceptions:</h3>
|
||
<pre><code> - SyntaxError: If the script contains a syntax error, it is caught and
|
||
returned as a part of the error message.
|
||
</code></pre></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
<dt id="connpy.ai"><code class="flex name class">
|
||
<span>class <span class="ident">ai</span></span>
|
||
<span>(</span><span>config,<br>org=None,<br>api_key=None,<br>engineer_model=None,<br>architect_model=None,<br>engineer_api_key=None,<br>architect_api_key=None,<br>console=None,<br>confirm_handler=None,<br>trust=False)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@ClassHook
|
||
class ai:
|
||
"""Hybrid Multi-Agent System: Selective Escalation with Role Persistence."""
|
||
|
||
SAFE_COMMANDS = [
|
||
r'^show\s+', r'^ls\s*', r'^cat\s+', r'^ip\s+', r'^pwd$', r'^hostname$', r'^uname',
|
||
r'^df\s*', r'^free\s*', r'^ps\s*', r'^ping\s+', r'^traceroute\s+', r'^whois\s+',
|
||
r'^kubectl\s+(get|describe|version|logs|top|explain|cluster-info|api-resources|api-versions)\s+',
|
||
r'^systemctl\s+status\s+', r'^journalctl\s+'
|
||
]
|
||
|
||
def __init__(self, config, org=None, api_key=None, engineer_model=None, architect_model=None, engineer_api_key=None, architect_api_key=None, console=None, confirm_handler=None, trust=False):
|
||
self.config = config
|
||
self.console = console or printer.console
|
||
self.confirm_handler = confirm_handler or self._local_confirm_handler
|
||
self.trusted_session = trust # Trust mode for the entire session
|
||
self.interrupted = False
|
||
|
||
|
||
# 1. Cargar configuración genérica
|
||
aiconfig = self.config.config.get("ai", {})
|
||
|
||
# Modelos (Prioridad: Argumento -> Config -> Default)
|
||
self.engineer_model = engineer_model or aiconfig.get("engineer_model") or "gemini/gemini-3.1-flash-lite"
|
||
self.architect_model = architect_model or aiconfig.get("architect_model") or "anthropic/claude-sonnet-4-6"
|
||
|
||
# API Keys (Prioridad: Argumento -> Config)
|
||
self.engineer_key = engineer_api_key or aiconfig.get("engineer_api_key")
|
||
self.architect_key = architect_api_key or aiconfig.get("architect_api_key")
|
||
|
||
# Custom Trusted Commands Regexes
|
||
custom_trusted = aiconfig.get("trusted_commands", [])
|
||
if isinstance(custom_trusted, str):
|
||
custom_trusted = [c.strip() for c in custom_trusted.split(",") if c.strip()]
|
||
self.safe_commands = list(self.SAFE_COMMANDS) + (custom_trusted if isinstance(custom_trusted, list) else [])
|
||
|
||
# Límites
|
||
self.max_history = 30
|
||
self.max_truncate = 50000
|
||
self.soft_limit_iterations = 20 # Show warning and suggest Ctrl+C
|
||
self.hard_limit_iterations = 50 # Force stop
|
||
|
||
# External tool registry (populated by plugins via ClassHook.modify)
|
||
self.external_engineer_tools = [] # Tool defs for Engineer LLM
|
||
self.external_architect_tools = [] # Tool defs for Architect LLM
|
||
self.external_tool_handlers = {} # {"tool_name": handler_callable}
|
||
self.tool_status_formatters = {} # {"tool_name": formatter_callable}
|
||
self.engineer_prompt_extensions = [] # Extra text for engineer prompt
|
||
self.architect_prompt_extensions = [] # Extra text for architect prompt
|
||
|
||
# MCP Manager
|
||
self.mcp_manager = MCPClientManager(self.config)
|
||
|
||
# Long-term memory
|
||
self.memory_path = os.path.join(self.config.defaultdir, "ai_memory.md")
|
||
self.long_term_memory = ""
|
||
if os.path.exists(self.memory_path):
|
||
try:
|
||
with open(self.memory_path, "r") as f:
|
||
self.long_term_memory = f.read()
|
||
except FileNotFoundError:
|
||
self.long_term_memory = ""
|
||
except PermissionError as e:
|
||
self.console.print(f"[warning]Warning: Cannot read AI memory file: {e}[/warning]")
|
||
except Exception as e:
|
||
self.console.print(f"[warning]Warning: Failed to load AI memory: {e}[/warning]")
|
||
|
||
# Session Management
|
||
self.sessions_dir = os.path.join(self.config.defaultdir, "ai_sessions")
|
||
os.makedirs(self.sessions_dir, exist_ok=True)
|
||
self.session_id = None
|
||
self.session_path = None
|
||
|
||
# Prompts base agnósticos
|
||
architect_instructions = ""
|
||
if self.architect_key:
|
||
architect_instructions = """
|
||
CRITICAL - CONSULT vs ESCALATE:
|
||
- ALWAYS use 'consult_architect' for: Configuration planning, design decisions, complex troubleshooting.
|
||
Examples: "consultalo con el arquitecto", "preguntale al arquitecto", "que opina el arquitecto"
|
||
You stay in control and present the advice to the user.
|
||
|
||
- ONLY use 'escalate_to_architect' when user EXPLICITLY asks to TALK to the Architect:
|
||
Examples: "quiero hablar con el arquitecto", "pasame con el arquitecto", "que me atienda el arquitecto"
|
||
After escalation, you hand over control completely.
|
||
|
||
- DEFAULT: When in doubt, use 'consult_architect'. Escalation is rare.
|
||
"""
|
||
else:
|
||
architect_instructions = """
|
||
CRITICAL - ARCHITECT UNAVAILABLE:
|
||
- The Strategic Reasoning Engine (Architect) is currently UNAVAILABLE because its API key is not configured.
|
||
- DO NOT attempt to consult or escalate to the architect.
|
||
- If the user asks to consult the architect, inform them that the Architect is offline and offer to help them directly to the best of your abilities.
|
||
"""
|
||
|
||
self._engineer_base_prompt = dedent(f"""
|
||
Role: TECHNICAL EXECUTION ENGINE.
|
||
Expertise: Universal Networking (Cisco, Nokia, Juniper, 6wind, etc.).
|
||
|
||
Rules:
|
||
- BE FAST AND EXTREMELY CONCISE: Provide direct answers. No filler words, no decorative language, no polite pleasantries. Save output tokens at all costs.
|
||
- KNOWLEDGE FIRST: For general networking questions (AS numbers, protocol details, standards, generic commands), use your internal knowledge. ONLY use tools when the user's specific infrastructure data is required.
|
||
- INVENTORY ONLY: 'run_commands', 'list_nodes', and 'get_node_info' are ONLY for interacting with the user's inventory.
|
||
- BROADCAST RESTRICTION: Avoid using filter '.*' in 'run_commands' unless the user explicitly requests a global action. Try to target specific nodes or groups based on the conversation.
|
||
- AUTONOMY: Proactively use iterative tool calls to find the root cause of infrastructure issues.
|
||
- BATCH OPERATIONS: When working on multiple devices, call tools in parallel.
|
||
- COMPLETE MISSIONS: Execute ALL steps of a mission before reporting back.
|
||
- DIAGRAM: Use ASCII art or Unicode box-drawing characters directly in your responses to visualize topologies or paths when helpful.
|
||
- EVIDENCE: Include 'Key Snippets' from tool outputs. Be token-efficient.
|
||
- NO WANDERING: Do not speculate. If stuck, report attempts.
|
||
- SAFETY: When you use 'run_commands' with configuration commands, the system automatically prompts the user for confirmation. Just execute - don't ask permission first.
|
||
{architect_instructions}
|
||
Network Context: {{self.long_term_memory if self.long_term_memory else "Empty."}}
|
||
""").strip()
|
||
|
||
self._architect_base_prompt = dedent(f"""
|
||
Role: STRATEGIC REASONING ENGINE.
|
||
Expertise: Network Architecture, Complex Troubleshooting, and Design Validation.
|
||
|
||
Rules:
|
||
- CONCISENESS IS MANDATORY: Strip out fluff, decorative language, and filler words. Provide direct, tactical instructions and analysis to save output tokens.
|
||
- STRATEGY: Define technical missions for the Engineer.
|
||
- DIAGRAM: Use ASCII art or Unicode box-drawing characters in your responses to visualize topologies, traffic paths, or logic flows.
|
||
- ENGINEER CAPABILITIES: Your Engineer can:
|
||
* Filter nodes (list_nodes), Run CLI commands (run_commands), Get metadata (get_node_info).
|
||
- ANALYSIS: Review technical findings to identify patterns or design failures.
|
||
- MEMORY: Update long-term facts ONLY when the user explicitly requests it.
|
||
|
||
CRITICAL - EFFICIENT DELEGATION:
|
||
- Plan ALL tasks upfront before delegating.
|
||
- Delegate ONCE with a complete, detailed mission including ALL steps.
|
||
- Example: "List all routers matching 'border.*', then run 'show ip bgp summary' and 'show ip route' on each, then analyze the outputs."
|
||
- DO NOT delegate multiple times for the same goal. Batch everything into ONE mission.
|
||
- Wait for Engineer's complete report before responding to user.
|
||
|
||
CRITICAL - RETURNING CONTROL:
|
||
- When your strategic analysis is complete and no further architectural decisions are needed, use 'return_to_engineer' to hand control back.
|
||
- The Engineer is better suited for ongoing technical execution and troubleshooting.
|
||
- Only stay in control if the user explicitly needs strategic oversight for multiple interactions.
|
||
|
||
Network Context: {self.long_term_memory if self.long_term_memory else "Empty."}
|
||
""").strip()
|
||
|
||
def _local_confirm_handler(self, prompt, default="n"):
|
||
"""Default confirmation handler using rich.prompt."""
|
||
from rich.prompt import Prompt
|
||
return Prompt.ask(prompt, default=default)
|
||
|
||
@property
|
||
def engineer_system_prompt(self):
|
||
"""Build engineer system prompt with plugin extensions."""
|
||
if self.engineer_prompt_extensions:
|
||
extensions = "\n".join(self.engineer_prompt_extensions)
|
||
return self._engineer_base_prompt + f"\n\nPlugin Capabilities:\n{extensions}"
|
||
return self._engineer_base_prompt
|
||
|
||
@property
|
||
def architect_system_prompt(self):
|
||
"""Build architect system prompt with plugin extensions."""
|
||
if self.architect_prompt_extensions:
|
||
extensions = "\n".join(self.architect_prompt_extensions)
|
||
return self._architect_base_prompt + f"\n\nPlugin Capabilities:\n{extensions}"
|
||
return self._architect_base_prompt
|
||
|
||
def register_ai_tool(self, tool_definition, handler, target="engineer", engineer_prompt=None, architect_prompt=None, status_formatter=None):
|
||
"""Register an external tool for the AI system.
|
||
|
||
Args:
|
||
tool_definition (dict): OpenAI-compatible tool definition.
|
||
handler (callable): Function(ai_instance, **tool_args) -> str.
|
||
target (str): 'engineer', 'architect', or 'both'.
|
||
engineer_prompt (str): Extra text for engineer system prompt.
|
||
architect_prompt (str): Extra text for architect system prompt.
|
||
status_formatter (callable): Function(args_dict) -> status string.
|
||
"""
|
||
name = tool_definition["function"]["name"]
|
||
|
||
# Check if already registered to prevent duplicates
|
||
if target in ("engineer", "both"):
|
||
if not any(t["function"]["name"] == name for t in self.external_engineer_tools):
|
||
self.external_engineer_tools.append(tool_definition)
|
||
if target in ("architect", "both"):
|
||
if not any(t["function"]["name"] == name for t in self.external_architect_tools):
|
||
self.external_architect_tools.append(tool_definition)
|
||
|
||
self.external_tool_handlers[name] = handler
|
||
|
||
if engineer_prompt and engineer_prompt not in self.engineer_prompt_extensions:
|
||
self.engineer_prompt_extensions.append(engineer_prompt)
|
||
if architect_prompt and architect_prompt not in self.architect_prompt_extensions:
|
||
self.architect_prompt_extensions.append(architect_prompt)
|
||
if status_formatter:
|
||
self.tool_status_formatters[name] = status_formatter
|
||
|
||
def _stream_completion(self, model, messages, tools, api_key, status=None, label="", debug=False, chunk_callback=None, **kwargs):
|
||
"""Stream a completion call, rendering styled Markdown in real-time.
|
||
|
||
Returns (response, streamed) where:
|
||
- response: reconstructed ModelResponse (same as non-streaming)
|
||
- streamed: True if text was rendered to console during streaming
|
||
"""
|
||
from rich.live import Live
|
||
|
||
stream_resp = completion(model=model, messages=messages, tools=tools, api_key=api_key, stream=True, **kwargs)
|
||
|
||
chunks = []
|
||
full_content = ""
|
||
is_streaming_text = False
|
||
has_tool_calls = False
|
||
live_display = None
|
||
|
||
# Determine styling based on current brain
|
||
role_label = "Network Architect" if "architect" in label.lower() else "Network Engineer"
|
||
alias = "architect" if "architect" in label.lower() else "engineer"
|
||
title = f"[bold {alias}]{role_label}[/bold {alias}]"
|
||
border = alias
|
||
|
||
try:
|
||
for chunk in stream_resp:
|
||
chunks.append(chunk)
|
||
delta = chunk.choices[0].delta
|
||
|
||
# Detect tool calls
|
||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||
has_tool_calls = True
|
||
|
||
# Stream text content with styled rendering
|
||
if hasattr(delta, 'content') and delta.content:
|
||
full_content += delta.content
|
||
|
||
if chunk and chunk_callback:
|
||
# Check for remote interruption during streaming
|
||
if hasattr(self, "interrupted") and self.interrupted:
|
||
raise KeyboardInterrupt
|
||
chunk_callback(delta.content)
|
||
|
||
if not chunk_callback:
|
||
if not is_streaming_text:
|
||
# Stop spinner definitively
|
||
if status:
|
||
try:
|
||
status.stop()
|
||
except Exception:
|
||
pass
|
||
|
||
# Create a stable, direct Console to bypass _ConsoleProxy recreation bugs
|
||
from rich.console import Console as RichConsole
|
||
from .printer import connpy_theme, get_original_stdout
|
||
stable_console = RichConsole(theme=connpy_theme, file=get_original_stdout())
|
||
|
||
live_display = Live(
|
||
Panel(Markdown(full_content), title=title, border_style=border, expand=False),
|
||
console=stable_console,
|
||
refresh_per_second=8,
|
||
transient=False
|
||
)
|
||
live_display.start()
|
||
is_streaming_text = True
|
||
else:
|
||
live_display.update(
|
||
Panel(Markdown(full_content), title=title, border_style=border, expand=False)
|
||
)
|
||
except Exception as e:
|
||
if not chunks:
|
||
raise
|
||
finally:
|
||
if live_display:
|
||
# Render final state with complete content
|
||
try:
|
||
live_display.update(
|
||
Panel(Markdown(full_content), title=title, border_style=border, expand=False)
|
||
)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
live_display.stop()
|
||
except Exception:
|
||
pass
|
||
|
||
# Rebuild complete response from chunks
|
||
try:
|
||
response = stream_chunk_builder(chunks, messages=messages)
|
||
except Exception:
|
||
# Fallback: manual reconstruction if stream_chunk_builder fails
|
||
full_content_rebuilt = ""
|
||
tool_calls_map = {}
|
||
for c in chunks:
|
||
d = c.choices[0].delta
|
||
if hasattr(d, 'content') and d.content:
|
||
full_content_rebuilt += d.content
|
||
if hasattr(d, 'tool_calls') and d.tool_calls:
|
||
for tc in d.tool_calls:
|
||
idx = tc.index
|
||
if idx not in tool_calls_map:
|
||
tool_calls_map[idx] = {"id": tc.id or "", "type": "function", "function": {"name": getattr(tc.function, 'name', '') or '', "arguments": getattr(tc.function, 'arguments', '') or ''}}
|
||
else:
|
||
if tc.id: tool_calls_map[idx]["id"] = tc.id
|
||
if tc.function:
|
||
if tc.function.name: tool_calls_map[idx]["function"]["name"] = tc.function.name
|
||
if tc.function.arguments: tool_calls_map[idx]["function"]["arguments"] += tc.function.arguments
|
||
|
||
# Build a minimal response-like object
|
||
class FakeFunc:
|
||
def __init__(self, name, arguments): self.name = name; self.arguments = arguments
|
||
class FakeTC:
|
||
def __init__(self, d): self.id = d["id"]; self.function = FakeFunc(d["function"]["name"], d["function"]["arguments"])
|
||
def model_dump(self, **kw): return {"id": self.id, "type": "function", "function": {"name": self.function.name, "arguments": self.function.arguments}}
|
||
class FakeMsg:
|
||
def __init__(self, content, tcs): self.content = content or None; self.tool_calls = tcs if tcs else None; self.role = "assistant"
|
||
def model_dump(self, **kw):
|
||
d = {"role": "assistant", "content": self.content}
|
||
if self.tool_calls: d["tool_calls"] = [tc.model_dump() for tc in self.tool_calls]
|
||
return d
|
||
class FakeChoice:
|
||
def __init__(self, msg): self.message = msg
|
||
class FakeResp:
|
||
def __init__(self, choice): self.choices = [choice]; self.usage = None
|
||
|
||
tcs = [FakeTC(tool_calls_map[i]) for i in sorted(tool_calls_map)] if tool_calls_map else None
|
||
response = FakeResp(FakeChoice(FakeMsg(full_content_rebuilt or full_content, tcs)))
|
||
|
||
# Only count as "streamed" if we rendered text AND it was the final response (no tool calls)
|
||
streamed = is_streaming_text and not has_tool_calls
|
||
return response, streamed
|
||
|
||
def _sanitize_messages(self, messages):
|
||
"""Sanitize message list for strict providers like Gemini.
|
||
|
||
Ensures that:
|
||
1. Every assistant message with tool_calls is followed by ALL its tool responses
|
||
2. No user/system messages appear between tool_calls and tool responses
|
||
3. Orphaned tool_calls at the end are removed
|
||
4. Orphaned tool responses without a preceding tool_call are removed
|
||
5. Incompatible metadata like cache_control is stripped for non-Anthropic models
|
||
6. Enforces strict alternating history to prevent BadRequestError on Gemini.
|
||
"""
|
||
if not messages:
|
||
return messages
|
||
|
||
# Pre-process messages to pull text from list contents (Anthropic cache format)
|
||
# and remove explicit cache keys.
|
||
pre_sanitized = []
|
||
for msg in messages:
|
||
m = msg.copy() if isinstance(msg, dict) else msg.model_dump(exclude_none=True)
|
||
|
||
# Convert content list to plain string if it's a system message with caching metadata
|
||
if m.get('role') == 'system' and isinstance(m.get('content'), list):
|
||
if m['content'] and isinstance(m['content'][0], dict) and m['content'][0].get('text'):
|
||
m['content'] = m['content'][0]['text']
|
||
else:
|
||
m['content'] = ""
|
||
|
||
# Remove any explicit cache_control key anywhere
|
||
if 'cache_control' in m: del m['cache_control']
|
||
if isinstance(m.get('content'), list):
|
||
for item in m['content']:
|
||
if isinstance(item, dict) and 'cache_control' in item: del item['cache_control']
|
||
|
||
pre_sanitized.append(m)
|
||
|
||
sanitized = []
|
||
last_role = None
|
||
|
||
i = 0
|
||
while i < len(pre_sanitized):
|
||
msg = pre_sanitized[i]
|
||
role = msg.get('role', '')
|
||
|
||
if role == 'system':
|
||
sanitized.append(msg)
|
||
last_role = 'system'
|
||
i += 1
|
||
|
||
elif role == 'user':
|
||
if last_role == 'user' and sanitized:
|
||
# Combine consecutive user messages
|
||
sanitized[-1]['content'] = str(sanitized[-1].get('content', '') or '') + '\n' + str(msg.get('content', '') or '')
|
||
else:
|
||
sanitized.append(msg)
|
||
last_role = 'user'
|
||
i += 1
|
||
|
||
elif role == 'assistant':
|
||
has_tools = bool(msg.get('tool_calls'))
|
||
|
||
# Gemini strict sequence: Assistant MUST be preceded by user or tool.
|
||
# If preceded by system, assistant, or if it's the very first message...
|
||
if last_role not in ('user', 'tool'):
|
||
sanitized.append({"role": "user", "content": "[System sequence separator: History Truncated/Merged]"})
|
||
last_role = 'user'
|
||
|
||
if has_tools:
|
||
# Look ahead for matching tool responses
|
||
tool_responses = []
|
||
j = i + 1
|
||
while j < len(pre_sanitized):
|
||
next_msg = pre_sanitized[j]
|
||
if next_msg.get('role') == 'tool':
|
||
tool_responses.append(next_msg)
|
||
j += 1
|
||
else:
|
||
break
|
||
|
||
if tool_responses:
|
||
sanitized.append(msg)
|
||
sanitized.extend(tool_responses)
|
||
last_role = 'tool'
|
||
i = j
|
||
else:
|
||
# Orphaned tool_calls with no responses - skip the assistant message
|
||
# If we just added a dummy user message for this assistant, remove it too
|
||
if sanitized and sanitized[-1].get('content') == "[System sequence separator: History Truncated/Merged]":
|
||
sanitized.pop()
|
||
last_role = sanitized[-1].get('role', '') if sanitized else None
|
||
i += 1
|
||
else:
|
||
sanitized.append(msg)
|
||
last_role = 'assistant'
|
||
i += 1
|
||
|
||
elif role == 'tool':
|
||
# Orphaned tool response (no preceding assistant with tool_calls) - skip
|
||
i += 1
|
||
|
||
else:
|
||
sanitized.append(msg)
|
||
last_role = role
|
||
i += 1
|
||
|
||
return sanitized
|
||
|
||
def _truncate(self, text, limit=None):
|
||
"""Truncate text to specified limit, keeping head (60%) and tail (40%)."""
|
||
if not isinstance(text, str): return str(text)
|
||
final_limit = limit or self.max_truncate
|
||
if len(text) <= final_limit: return text
|
||
head_limit = int(final_limit * 0.6)
|
||
tail_limit = int(final_limit * 0.4)
|
||
return (text[:head_limit] + f"\n\n[... OUTPUT TRUNCATED ...]\n\n" + text[-tail_limit:])
|
||
|
||
def _print_debug_observation(self, fn, obs, status=None):
|
||
"""Prints a tool observation in a readable way during debug mode."""
|
||
# Try to parse as JSON if it's a string
|
||
if isinstance(obs, str):
|
||
try:
|
||
obs_data = json.loads(obs)
|
||
except Exception:
|
||
obs_data = obs
|
||
else:
|
||
obs_data = obs
|
||
|
||
if isinstance(obs_data, dict):
|
||
elements = []
|
||
for k, v in obs_data.items():
|
||
elements.append(Text(f"• {k}:", style="key"))
|
||
# Use Text for values to ensure newlines are rendered
|
||
val = str(v)
|
||
# If it's a multiline string from a delegation task, keep it clean
|
||
elements.append(Text(val))
|
||
|
||
if not elements:
|
||
content = Text("Empty data set")
|
||
else:
|
||
# Add a small spacer instead of a Rule for cleaner look
|
||
from rich.console import Group
|
||
content = Group(*elements)
|
||
elif isinstance(obs_data, list):
|
||
content = Text("\n".join(f"• {item}" for item in obs_data))
|
||
else:
|
||
content = Text(str(obs_data))
|
||
|
||
title = f"[bold]{fn}[/bold]"
|
||
|
||
# Stop status before printing panel to avoid ghosting
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
|
||
self.console.print(Panel(content, title=title, border_style="ai_status"))
|
||
|
||
# Resume status
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
def manage_memory_tool(self, content, action="append"):
|
||
"""Save or update long-term memory. Only use when user explicitly requests it."""
|
||
if not content or not content.strip():
|
||
return "Error: Cannot save empty content to memory."
|
||
|
||
try:
|
||
mode = "a" if action == "append" else "w"
|
||
os.makedirs(os.path.dirname(self.memory_path), exist_ok=True)
|
||
with open(self.memory_path, mode) as f:
|
||
timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M')
|
||
f.write(f"\n\n## {timestamp}\n{content.strip()}\n" if action == "append" else content)
|
||
|
||
# Reload memory after update
|
||
with open(self.memory_path, "r") as f:
|
||
self.long_term_memory = f.read()
|
||
|
||
return "Memory updated successfully."
|
||
except PermissionError as e:
|
||
return f"Error: Permission denied writing to memory file: {e}"
|
||
except Exception as e:
|
||
return f"Error updating memory: {str(e)}"
|
||
|
||
|
||
def list_nodes_tool(self, filter_pattern=".*"):
|
||
"""List nodes matching the filter pattern. Returns metadata for <=5 nodes, names only for more."""
|
||
try:
|
||
matched_names = self.config._getallnodes(filter_pattern)
|
||
if not matched_names: return "No nodes found."
|
||
if len(matched_names) <= 5:
|
||
matched_data = self.config.getitems(matched_names, extract=True)
|
||
res = {}
|
||
for name, data in matched_data.items():
|
||
os_tag = "unknown"
|
||
if isinstance(data, dict):
|
||
ts = data.get("tags")
|
||
if isinstance(ts, dict): os_tag = ts.get("os", "unknown")
|
||
res[name] = {"os": os_tag}
|
||
return res
|
||
return {"count": len(matched_names), "nodes": matched_names, "note": "Use 'get_node_info' for details."}
|
||
except Exception as e:
|
||
return f"Error listing nodes: {str(e)}"
|
||
|
||
def _is_safe_command(self, cmd):
|
||
"""Check if a command matches safe patterns."""
|
||
return any(re.match(pattern, cmd.strip(), re.IGNORECASE) for pattern in self.safe_commands)
|
||
|
||
def run_commands_tool(self, nodes_filter, commands, status=None):
|
||
"""Execute commands on nodes matching the filter. Native interactive confirmation for unsafe commands."""
|
||
# Handle if commands is a JSON string
|
||
if isinstance(commands, str):
|
||
try:
|
||
commands = json.loads(commands)
|
||
except ValueError:
|
||
commands = [c.strip() for c in commands.split('\n') if c.strip()]
|
||
|
||
# Expand multi-line commands within a list (in case the AI packs them)
|
||
if isinstance(commands, list):
|
||
expanded_commands = []
|
||
for cmd in commands:
|
||
expanded_commands.extend([c.strip() for c in str(cmd).split('\n') if c.strip()])
|
||
commands = expanded_commands
|
||
else:
|
||
commands = [str(commands)]
|
||
|
||
# Check command safety natively
|
||
if not self.trusted_session:
|
||
unsafe_commands = [cmd for cmd in commands if not self._is_safe_command(cmd)]
|
||
if unsafe_commands:
|
||
# Stop the spinner so prompt doesn't get messed up
|
||
if status: status.stop()
|
||
|
||
# Show ALL commands with unsafe ones highlighted
|
||
formatted_cmds = []
|
||
for cmd in commands:
|
||
if cmd in unsafe_commands:
|
||
formatted_cmds.append(f" • [warning]{cmd}[/warning]")
|
||
else:
|
||
formatted_cmds.append(f" • {cmd}")
|
||
|
||
panel_content = f"Target: {nodes_filter}\nCommands:\n" + "\n".join(formatted_cmds)
|
||
# Use print_important if available (for remote bridges) fallback to standard print
|
||
print_fn = getattr(self.console, "print_important", self.console.print)
|
||
print_fn(Panel(panel_content, title="[bold warning]⚠️ UNSAFE COMMANDS DETECTED[/bold warning]", border_style="warning"))
|
||
|
||
try:
|
||
user_resp = self.confirm_handler("[bold warning]Execute? (y: yes / n: no / a: allow all this session / <text>: feedback)[/bold warning]", default="n")
|
||
except KeyboardInterrupt:
|
||
if status: status.update("[ai_status]Engineer: Resuming...")
|
||
self.console.print("[fail]✗ Aborted by user (Ctrl+C).[/fail]")
|
||
raise
|
||
|
||
# Resume the spinner
|
||
if status: status.update("[ai_status]Engineer: Processing user response...")
|
||
|
||
user_resp_lower = user_resp.strip().lower()
|
||
if user_resp_lower in ['a', 'allow']:
|
||
self.trusted_session = True
|
||
self.console.print("[pass]✓ Trust Mode Enabled. All future commands in this session will execute without confirmation.[/pass]")
|
||
elif user_resp_lower in ['y', 'yes']:
|
||
self.console.print("[pass]✓ Executing...[/pass]")
|
||
elif user_resp_lower in ['n', 'no', '', 'cancel']:
|
||
self.console.print("[fail]✗ Execution rejected by user.[/fail]")
|
||
return "Error: User rejected execution."
|
||
else:
|
||
self.console.print(f"[user_prompt]User feedback: [/user_prompt]{user_resp}")
|
||
return f"User requested changes: {user_resp}. Please adjust the commands based on this feedback and try again."
|
||
|
||
try:
|
||
matched_names = self.config._getallnodes(nodes_filter)
|
||
if not matched_names: return "No nodes found matching filter."
|
||
thisnodes_dict = self.config.getitems(matched_names, extract=True)
|
||
result = nodes(thisnodes_dict, config=self.config).run(commands)
|
||
return result
|
||
except Exception as e:
|
||
return f"Error executing commands: {str(e)}"
|
||
|
||
def get_node_info_tool(self, node_name):
|
||
"""Get detailed metadata for a specific node. Passwords are masked."""
|
||
try:
|
||
d = self.config.getitem(node_name, extract=True)
|
||
if 'password' in d: d['password'] = '***'
|
||
return d
|
||
except Exception as e:
|
||
return f"Error getting node info: {str(e)}"
|
||
|
||
def _engineer_loop(self, task, status=None, debug=False, chat_history=None):
|
||
"""Internal loop where the Engineer executes technical tasks for the Architect."""
|
||
# Optimización de caché para el Ingeniero (Solo para Anthropic directo, Vertex tiene reglas distintas)
|
||
if "claude" in self.engineer_model.lower() and "vertex" not in self.engineer_model.lower():
|
||
messages = [{"role": "system", "content": [{"type": "text", "text": self.engineer_system_prompt, "cache_control": {"type": "ephemeral"}}]}]
|
||
else:
|
||
messages = [{"role": "system", "content": self.engineer_system_prompt}]
|
||
|
||
if chat_history:
|
||
# Clean chat history from caching metadata if engineer is not a compatible Claude model
|
||
if "claude" not in self.engineer_model.lower() or "vertex" in self.engineer_model.lower():
|
||
messages.extend(self._sanitize_messages(chat_history[-5:]))
|
||
else:
|
||
messages.extend(chat_history[-5:])
|
||
|
||
messages.append({"role": "user", "content": f"MISSION: {task}"})
|
||
|
||
tools = self._get_engineer_tools()
|
||
usage = {"input": 0, "output": 0, "total": 0}
|
||
iteration = 0
|
||
soft_limit_warned = False
|
||
|
||
try:
|
||
# Set up remote interrupt callback if bridge is provided
|
||
if status and hasattr(status, "on_interrupt"):
|
||
status.on_interrupt = lambda: setattr(self, "interrupted", True)
|
||
|
||
while iteration < self.hard_limit_iterations:
|
||
iteration += 1
|
||
|
||
# Check for interruption
|
||
if self.interrupted:
|
||
raise KeyboardInterrupt
|
||
|
||
# Soft limit warning
|
||
if iteration == self.soft_limit_iterations and not soft_limit_warned:
|
||
self.console.print(f"[warning]⚠ Engineer has performed {iteration} steps. This is taking longer than expected.[/warning]")
|
||
self.console.print(f"[warning] You can press Ctrl+C to interrupt and get a summary.[/warning]")
|
||
soft_limit_warned = True
|
||
|
||
if status and not chat_history: status.update(f"[ai_status]Engineer: Analyzing mission... (step {iteration})")
|
||
|
||
try:
|
||
safe_messages = self._sanitize_messages(messages)
|
||
response = completion(model=self.engineer_model, messages=safe_messages, tools=tools, api_key=self.engineer_key)
|
||
except Exception as e:
|
||
if status: status.stop()
|
||
raise ValueError(f"Engineer failed to connect: {str(e)}")
|
||
|
||
if hasattr(response, "usage") and response.usage:
|
||
usage["input"] += getattr(response.usage, "prompt_tokens", 0)
|
||
usage["output"] += getattr(response.usage, "completion_tokens", 0)
|
||
usage["total"] += getattr(response.usage, "total_tokens", 0)
|
||
|
||
resp_msg = response.choices[0].message
|
||
msg_dict = resp_msg.model_dump(exclude_none=True)
|
||
if msg_dict.get("tool_calls") and msg_dict.get("content") == "": msg_dict["content"] = None
|
||
messages.append(msg_dict)
|
||
|
||
if not resp_msg.tool_calls: break
|
||
for tc in resp_msg.tool_calls:
|
||
fn, args = tc.function.name, json.loads(tc.function.arguments)
|
||
|
||
# Notificación en tiempo real de la tarea técnica (Only if not in Architect loop)
|
||
if status and not chat_history:
|
||
if fn == "list_nodes": status.update(f"[ai_status]Engineer: [SEARCH] {args.get('filter_pattern','.*')}")
|
||
elif fn == "run_commands":
|
||
cmds = args.get('commands', [])
|
||
cmd_str = cmds[0] if cmds else ""
|
||
status.update(f"[ai_status]Engineer: [CMD] {cmd_str}")
|
||
elif fn == "get_node_info": status.update(f"[ai_status]Engineer: [INSPECT] {args.get('node_name','')}")
|
||
elif fn.startswith("mcp_"):
|
||
server = fn.split("__")[0].replace("mcp_", "")
|
||
tool = fn.split("__")[1] if "__" in fn else fn
|
||
status.update(f"[ai_status]Engineer: [MCP:{server}] {tool}")
|
||
elif fn in self.tool_status_formatters: status.update(self.tool_status_formatters[fn](args))
|
||
|
||
if debug:
|
||
self._print_debug_observation(f"Decision: {fn}", args, status=status)
|
||
|
||
if fn == "list_nodes": obs = self.list_nodes_tool(**args)
|
||
elif fn == "run_commands": obs = self.run_commands_tool(**args, status=status)
|
||
elif fn == "get_node_info": obs = self.get_node_info_tool(**args)
|
||
elif fn.startswith("mcp_"):
|
||
obs = run_ai_async(self.mcp_manager.call_tool(fn, args)).result(timeout=60)
|
||
elif fn in self.external_tool_handlers: obs = self.external_tool_handlers[fn](self, **args)
|
||
else: obs = f"Error: Unknown tool '{fn}'."
|
||
|
||
if debug:
|
||
self._print_debug_observation(f"Observation: {fn}", obs, status=status)
|
||
|
||
# Ensure observation is a string and truncated for the LLM
|
||
obs_str = obs if isinstance(obs, str) else json.dumps(obs)
|
||
messages.append({"tool_call_id": tc.id, "role": "tool", "name": fn, "content": self._truncate(obs_str)})
|
||
|
||
if iteration >= self.hard_limit_iterations:
|
||
self.console.print(f"[error]⛔ Engineer reached hard limit ({self.hard_limit_iterations} steps). Forcing stop.[/error]")
|
||
|
||
if debug and resp_msg.content:
|
||
self.console.print(Panel(Text(resp_msg.content), title="[bold engineer]Engineer Final Report to Architect[/bold engineer]", border_style="engineer"))
|
||
|
||
return resp_msg.content, usage
|
||
except Exception as e:
|
||
return f"Engineer failed: {str(e)}", usage
|
||
|
||
def _get_engineer_tools(self, os_filter: str = None):
|
||
"""Define tools available to the Engineer."""
|
||
base_tools = [
|
||
{"type": "function", "function": {"name": "list_nodes", "description": "[Universal Platform] Lists available nodes in the inventory.", "parameters": {"type": "object", "properties": {"filter_pattern": {"type": "string", "description": "Regex to filter nodes (e.g. '.*', 'border.*')."}}}}},
|
||
{"type": "function", "function": {"name": "run_commands", "description": "[Universal Platform] Runs one or more commands on matched nodes. MANDATORY: You MUST call 'list_nodes' first to verify the target list.", "parameters": {"type": "object", "properties": {"nodes_filter": {"type": "string", "description": "Exact node name or verified filter pattern."}, "commands": {"type": "array", "items": {"type": "string"}, "description": "List of commands (e.g. ['show ip route', 'show int desc'])."}}, "required": ["nodes_filter", "commands"]}}},
|
||
{"type": "function", "function": {"name": "get_node_info", "description": "[Universal Platform] Gets full metadata for a specific node.", "parameters": {"type": "object", "properties": {"node_name": {"type": "string"}}, "required": ["node_name"]}}}
|
||
]
|
||
|
||
# Add dynamic tools from MCP
|
||
try:
|
||
mcp_tools = run_ai_async(self.mcp_manager.get_tools_for_llm(os_filter=os_filter)).result(timeout=10)
|
||
base_tools.extend(mcp_tools)
|
||
except Exception as e:
|
||
# Silently fail for LLM tools
|
||
pass
|
||
|
||
if self.architect_key:
|
||
base_tools.extend([
|
||
{"type": "function", "function": {"name": "consult_architect", "description": "Ask the Strategic Reasoning Engine for advice on complex design, architecture, or troubleshooting decisions. You remain in control and will present the response to the user. Use this for: configuration planning, design validation, complex troubleshooting.", "parameters": {"type": "object", "properties": {"question": {"type": "string", "description": "Strategic question or decision needed."}, "technical_summary": {"type": "string", "description": "Technical findings and context gathered so far."}}, "required": ["question", "technical_summary"]}}},
|
||
{"type": "function", "function": {"name": "escalate_to_architect", "description": "Transfer full control to the Strategic Reasoning Engine. Use ONLY when the user explicitly requests the Architect or when the problem requires strategic oversight beyond consultation. After escalation, the Architect takes over the conversation.", "parameters": {"type": "object", "properties": {"reason": {"type": "string", "description": "Why you're escalating (e.g. 'User requested Architect', 'Complex multi-site design needed')."}, "context": {"type": "string", "description": "Full context and findings to hand over."}}, "required": ["reason", "context"]}}}
|
||
])
|
||
|
||
# Deduplicate by name to prevent Gemini BadRequestError
|
||
all_tools = base_tools + self.external_engineer_tools
|
||
seen_names = set()
|
||
unique_tools = []
|
||
for t in all_tools:
|
||
name = t["function"]["name"]
|
||
if name not in seen_names:
|
||
unique_tools.append(t)
|
||
seen_names.add(name)
|
||
return unique_tools
|
||
|
||
def _get_architect_tools(self):
|
||
"""Define tools available to the Strategic Reasoning Engine."""
|
||
base_tools = [
|
||
{"type": "function", "function": {"name": "delegate_to_engineer", "description": "Delegates a technical mission to the Engineer.", "parameters": {"type": "object", "properties": {"task": {"type": "string", "description": "Detailed technical mission or goal."}}, "required": ["task"]}}},
|
||
{"type": "function", "function": {"name": "return_to_engineer", "description": "Return control to the Engineer. Use this when your strategic analysis is complete and the Engineer should handle the rest of the conversation.", "parameters": {"type": "object", "properties": {"summary": {"type": "string", "description": "Brief summary of your analysis to hand over to the Engineer."}}, "required": ["summary"]}}},
|
||
{"type": "function", "function": {"name": "manage_memory_tool", "description": "Saves information to long-term memory. MANDATORY: Only use this if the user explicitly asks to remember or save something.", "parameters": {"type": "object", "properties": {"content": {"type": "string"}, "action": {"type": "string", "enum": ["append", "replace"]}}, "required": ["content"]}}}
|
||
]
|
||
|
||
all_tools = base_tools + self.external_architect_tools
|
||
seen_names = set()
|
||
unique_tools = []
|
||
for t in all_tools:
|
||
name = t["function"]["name"]
|
||
if name not in seen_names:
|
||
unique_tools.append(t)
|
||
seen_names.add(name)
|
||
return unique_tools
|
||
|
||
def _get_sessions(self):
|
||
"""Returns a list of session metadata sorted by date."""
|
||
sessions = []
|
||
if not os.path.exists(self.sessions_dir):
|
||
return []
|
||
for f in os.listdir(self.sessions_dir):
|
||
if f.endswith(".json"):
|
||
path = os.path.join(self.sessions_dir, f)
|
||
try:
|
||
with open(path, "r") as fs:
|
||
data = json.load(fs)
|
||
sessions.append({
|
||
"id": f[:-5],
|
||
"title": data.get("title", "Untitled Session"),
|
||
"created_at": data.get("created_at", "Unknown"),
|
||
"model": data.get("model", "Unknown"),
|
||
"path": path
|
||
})
|
||
except Exception:
|
||
continue
|
||
return sorted(sessions, key=lambda x: x["created_at"], reverse=True)
|
||
|
||
def list_sessions(self):
|
||
"""Prints a list of sessions using printer.table."""
|
||
sessions = self._get_sessions()
|
||
if not sessions:
|
||
printer.info("No saved AI sessions found.")
|
||
return
|
||
|
||
columns = ["ID", "Title", "Created At", "Model"]
|
||
rows = [[s["id"], s["title"], s["created_at"], s["model"]] for s in sessions]
|
||
printer.table("AI Persisted Sessions", columns, rows)
|
||
|
||
def load_session_data(self, session_id):
|
||
"""Loads a session's raw data by ID."""
|
||
path = os.path.join(self.sessions_dir, f"{session_id}.json")
|
||
if os.path.exists(path):
|
||
try:
|
||
with open(path, "r") as f:
|
||
data = json.load(f)
|
||
self.session_id = session_id
|
||
self.session_path = path
|
||
return data
|
||
except Exception as e:
|
||
printer.error(f"Failed to load session {session_id}: {e}")
|
||
return None
|
||
|
||
def delete_session(self, session_id):
|
||
"""Deletes a session by ID."""
|
||
path = os.path.join(self.sessions_dir, f"{session_id}.json")
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
printer.success(f"Session {session_id} deleted.")
|
||
else:
|
||
printer.error(f"Session {session_id} not found.")
|
||
|
||
def get_last_session_id(self):
|
||
"""Returns the ID of the most recent session."""
|
||
sessions = self._get_sessions()
|
||
return sessions[0]["id"] if sessions else None
|
||
|
||
def _generate_session_id(self, query):
|
||
"""Generates a unique session ID based on timestamp."""
|
||
return datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
|
||
|
||
def save_session(self, history, title=None, model=None):
|
||
"""Saves current history to the session file."""
|
||
if not self.session_id:
|
||
# Generate ID from first user query if available
|
||
first_user_msg = next((m["content"] for m in history if m["role"] == "user"), "new-session")
|
||
self.session_id = self._generate_session_id(first_user_msg)
|
||
self.session_path = os.path.join(self.sessions_dir, f"{self.session_id}.json")
|
||
|
||
# If it's a new file, we might want to set a better title
|
||
if not os.path.exists(self.session_path) and not title:
|
||
raw_title = next((m["content"] for m in history if m["role"] == "user"), "New Session")
|
||
# Clean title: remove newlines, multiple spaces
|
||
clean_title = " ".join(raw_title.split())
|
||
if len(clean_title) > 40:
|
||
title = clean_title[:37].strip() + "..."
|
||
else:
|
||
title = clean_title
|
||
|
||
try:
|
||
# Read existing metadata if it exists
|
||
metadata = {}
|
||
if os.path.exists(self.session_path):
|
||
with open(self.session_path, "r") as f:
|
||
metadata = json.load(f)
|
||
|
||
metadata.update({
|
||
"id": self.session_id,
|
||
"title": title or metadata.get("title", "New Session"),
|
||
"created_at": metadata.get("created_at", datetime.datetime.now().isoformat()),
|
||
"updated_at": datetime.datetime.now().isoformat(),
|
||
"model": model or metadata.get("model", self.engineer_model),
|
||
"history": history
|
||
})
|
||
|
||
with open(self.session_path, "w") as f:
|
||
json.dump(metadata, f, indent=4)
|
||
except Exception as e:
|
||
printer.error(f"Failed to save session: {e}")
|
||
|
||
except Exception as e:
|
||
printer.error(f"Failed to save session: {e}")
|
||
|
||
@MethodHook
|
||
def ask(self, user_input, dryrun=False, chat_history=None, status=None, debug=False, stream=True, session_id=None, chunk_callback=None):
|
||
if not self.engineer_key:
|
||
raise ValueError("Engineer API key not configured. Use 'connpy config --engineer-api-key <key>' to set it.")
|
||
|
||
if chat_history is None: chat_history = []
|
||
|
||
# Load session if provided and history is empty
|
||
if session_id and not chat_history:
|
||
session_data = self.load_session_data(session_id)
|
||
if session_data:
|
||
chat_history = session_data.get("history", [])
|
||
# If we loaded history, the caller might need it back
|
||
# But typically ask() is called in a loop with an external history object
|
||
|
||
usage = {"input": 0, "output": 0, "total": 0}
|
||
|
||
# 1. Selector de Rol inicial (Sticky Brain)
|
||
explicit_architect = re.match(r'^(architect|arquitecto|@architect)[:\s]', user_input, re.I)
|
||
explicit_engineer = re.match(r'^(engineer|ingeniero|@engineer)[:\s]', user_input, re.I)
|
||
|
||
if explicit_architect:
|
||
current_brain = "architect"
|
||
elif explicit_engineer:
|
||
current_brain = "engineer"
|
||
else:
|
||
# Sticky Brain: Detectar si el Arquitecto estaba al mando en el historial reciente
|
||
is_architect_active = False
|
||
for msg in reversed(chat_history[-5:]):
|
||
tcs = msg.get('tool_calls') if isinstance(msg, dict) else getattr(msg, 'tool_calls', None)
|
||
if tcs:
|
||
for tc in tcs:
|
||
fn = tc.get('function', {}).get('name') if isinstance(tc, dict) else getattr(getattr(tc, 'function', None), 'name', '')
|
||
# Architect stays in control if delegating tasks or if Engineer escalated to them
|
||
# consult_architect is just Engineer asking for advice - Engineer keeps control
|
||
if fn in ['delegate_to_engineer', 'escalate_to_architect']:
|
||
is_architect_active = True; break
|
||
if is_architect_active: break
|
||
current_brain = "architect" if is_architect_active else "engineer"
|
||
|
||
# 2. Preparación de mensajes y limpieza
|
||
clean_input = re.sub(r'^(architect|arquitecto|engineer|ingeniero|@architect|@engineer)[:\s]+', '', user_input, flags=re.IGNORECASE).strip()
|
||
|
||
system_prompt = self.architect_system_prompt if current_brain == "architect" else self.engineer_system_prompt
|
||
tools = self._get_architect_tools() if current_brain == "architect" else self._get_engineer_tools()
|
||
model = self.architect_model if current_brain == "architect" else self.engineer_model
|
||
key = self.architect_key if current_brain == "architect" else self.engineer_key
|
||
|
||
# Estructura optimizada para Prompt Caching (Solo para Anthropic directo, Vertex tiene reglas distintas)
|
||
if "claude" in model.lower() and "vertex" not in model.lower():
|
||
messages = [{"role": "system", "content": [{"type": "text", "text": system_prompt, "cache_control": {"type": "ephemeral"}}]}]
|
||
else:
|
||
messages = [{"role": "system", "content": system_prompt}]
|
||
|
||
# Interleaving de historial
|
||
last_role = "system"
|
||
# Sanitize history if the current target model is not compatible with cache_control
|
||
history_to_process = chat_history[-self.max_history:]
|
||
if "claude" not in model.lower() or "vertex" in model.lower():
|
||
history_to_process = self._sanitize_messages(history_to_process)
|
||
|
||
for msg in history_to_process:
|
||
m = msg if isinstance(msg, dict) else msg.model_dump(exclude_none=True)
|
||
role = m.get('role')
|
||
if role == last_role and role == 'user':
|
||
messages[-1]['content'] += "\n" + (m.get('content') or "")
|
||
continue
|
||
if role == 'assistant' and m.get('tool_calls') and m.get('content') == "": m['content'] = None
|
||
messages.append(m)
|
||
last_role = role
|
||
|
||
if last_role == 'user': messages[-1]['content'] += "\n" + clean_input
|
||
else: messages.append({"role": "user", "content": clean_input})
|
||
|
||
# 3. Bucle de ejecución
|
||
iteration = 0
|
||
try:
|
||
# Set up remote interrupt callback if bridge is provided
|
||
if status and hasattr(status, "on_interrupt"):
|
||
status.on_interrupt = lambda: setattr(self, "interrupted", True)
|
||
|
||
while iteration < self.hard_limit_iterations:
|
||
iteration += 1
|
||
|
||
# Check for interruption
|
||
if self.interrupted:
|
||
raise KeyboardInterrupt
|
||
|
||
# Soft limit warning
|
||
if iteration == self.soft_limit_iterations and not soft_limit_warned:
|
||
self.console.print(f"[warning]⚠ Agent has performed {iteration} steps. This is taking longer than expected.[/warning]")
|
||
self.console.print(f"[warning] You can press Ctrl+C to interrupt and get a summary of progress.[/warning]")
|
||
soft_limit_warned = True
|
||
|
||
label = "[architect][bold]Architect[/bold][/architect]" if current_brain == "architect" else "[engineer][bold]Engineer[/bold][/engineer]"
|
||
if status:
|
||
# Notify responder identity ONLY for web/remote clients (StatusBridge has is_web)
|
||
if getattr(status, "is_web", False):
|
||
status.update(f"__RESPONDER__:{current_brain}")
|
||
status.update(f"{label} is thinking... (step {iteration})")
|
||
|
||
streamed_response = False
|
||
try:
|
||
safe_messages = self._sanitize_messages(messages)
|
||
if stream:
|
||
response, streamed_response = self._stream_completion(
|
||
model=model, messages=safe_messages, tools=tools, api_key=key,
|
||
status=status, label=label, debug=debug, num_retries=3,
|
||
chunk_callback=chunk_callback
|
||
)
|
||
else:
|
||
response = completion(model=model, messages=safe_messages, tools=tools, api_key=key, num_retries=3)
|
||
except Exception as e:
|
||
if current_brain == "architect":
|
||
if status: status.update("[unavailable]Architect unavailable! Falling back to Engineer...")
|
||
# Preserve context when falling back - use clean_input directly
|
||
current_brain = "engineer"
|
||
model = self.engineer_model
|
||
tools = self._get_engineer_tools()
|
||
key = self.engineer_key
|
||
# Rebuild messages with Engineer system prompt and original user request
|
||
messages = [{"role": "system", "content": self.engineer_system_prompt}]
|
||
# Add chat history if exists (excluding system prompt)
|
||
if chat_history:
|
||
for msg in chat_history[-self.max_history:]:
|
||
if msg.get('role') != 'system':
|
||
messages.append(msg)
|
||
# Add current user request
|
||
messages.append({"role": "user", "content": clean_input})
|
||
continue
|
||
else:
|
||
return {"response": f"Error: Both engines failed. {str(e)}", "chat_history": messages[1:], "usage": usage}
|
||
|
||
if hasattr(response, "usage") and response.usage:
|
||
usage["input"] += getattr(response.usage, "prompt_tokens", 0)
|
||
usage["output"] += getattr(response.usage, "completion_tokens", 0)
|
||
usage["total"] += getattr(response.usage, "total_tokens", 0)
|
||
|
||
resp_msg = response.choices[0].message
|
||
msg_dict = resp_msg.model_dump(exclude_none=True)
|
||
if msg_dict.get("tool_calls") and msg_dict.get("content") == "": msg_dict["content"] = None
|
||
messages.append(msg_dict)
|
||
|
||
if debug and resp_msg.content and not streamed_response:
|
||
# In CLI debug mode, only print intermediate reasoning if there are tool calls AND it wasn't already streamed.
|
||
# If there are no tool calls, this content is the final answer and will be printed by the caller.
|
||
if resp_msg.tool_calls:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Markdown(resp_msg.content), title=f"[{current_brain}][bold]{label} Reasoning[/bold][/{current_brain}]", border_style="architect" if current_brain == "architect" else "engineer"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
if not resp_msg.tool_calls: break
|
||
|
||
# Track if we need to inject a user message after all tool responses
|
||
pending_user_message = None
|
||
|
||
for tc in resp_msg.tool_calls:
|
||
fn, args = tc.function.name, json.loads(tc.function.arguments)
|
||
|
||
# Validate tool access based on current brain
|
||
if fn in ['delegate_to_engineer'] and current_brain != "architect":
|
||
obs = f"Error: Tool '{fn}' is only available to the Architect (Architect). You are the Engineer (Engineer). Use 'run_commands' directly to execute configuration."
|
||
messages.append({"tool_call_id": tc.id, "role": "tool", "name": fn, "content": obs})
|
||
continue
|
||
|
||
if status:
|
||
if fn == "delegate_to_engineer": status.update(f"[architect]Architect: [DELEGATING MISSION] {args.get('task','')[:40]}...")
|
||
elif fn == "manage_memory_tool": status.update(f"[architect]Architect: [UPDATING MEMORY]")
|
||
|
||
if debug:
|
||
self._print_debug_observation(f"Decision: {fn}", args, status=status)
|
||
|
||
if fn == "delegate_to_engineer":
|
||
obs, eng_usage = self._engineer_loop(args["task"], status=status, debug=debug, chat_history=messages[:-1])
|
||
usage["input"] += eng_usage["input"]; usage["output"] += eng_usage["output"]; usage["total"] += eng_usage["total"]
|
||
elif fn == "consult_architect":
|
||
if status: status.update("[architect]Engineer consulting Architect...")
|
||
try:
|
||
# Consultation only - Engineer stays in control
|
||
claude_resp = completion(
|
||
model=self.architect_model,
|
||
messages=[
|
||
{"role": "system", "content": self.architect_system_prompt},
|
||
{"role": "user", "content": f"The Engineer needs your strategic advice.\n\nTECHNICAL SUMMARY: {args['technical_summary']}\n\nQUESTION: {args['question']}\n\nProvide strategic guidance. The Engineer will continue handling the user."}
|
||
],
|
||
api_key=self.architect_key,
|
||
num_retries=3
|
||
)
|
||
obs = claude_resp.choices[0].message.content
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Markdown(obs), title="[architect]Architect Consultation[/architect]", border_style="architect"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
except Exception as e:
|
||
if status: status.update("[unavailable]Architect unavailable! Engineer continuing alone...")
|
||
obs = f"Architect unavailable ({str(e)}). Proceeding with your best technical judgment."
|
||
|
||
elif fn == "escalate_to_architect":
|
||
if status: status.update("[architect]Transferring control to Architect...")
|
||
# Full escalation - Architect takes over
|
||
current_brain = "architect"
|
||
model = self.architect_model
|
||
tools = self._get_architect_tools()
|
||
key = self.architect_key
|
||
messages[0] = {"role": "system", "content": self.architect_system_prompt}
|
||
# Prepare handover context to inject AFTER all tool responses
|
||
handover_msg = f"HANDOVER FROM EXECUTION ENGINE\n\nReason: {args['reason']}\n\nContext: {args['context']}\n\nYou are now in control of this conversation."
|
||
pending_user_message = handover_msg
|
||
obs = "Control transferred to Architect. Handover context will be provided."
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Text(handover_msg), title="[architect]Escalation to Architect[/architect]", border_style="architect"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
elif fn == "return_to_engineer":
|
||
if status: status.update("[engineer]Transferring control back to Engineer...")
|
||
# Architect returns control to Engineer
|
||
current_brain = "engineer"
|
||
model = self.engineer_model
|
||
tools = self._get_engineer_tools()
|
||
key = self.engineer_key
|
||
messages[0] = {"role": "system", "content": self.engineer_system_prompt}
|
||
# Prepare handover context to inject AFTER all tool responses
|
||
handover_msg = f"HANDOVER FROM ARCHITECT\n\nSummary: {args['summary']}\n\nYou are now back in control. Continue handling the user's requests."
|
||
pending_user_message = handover_msg
|
||
obs = "Control returned to Engineer. Handover summary will be provided."
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Text(handover_msg), title="[engineer]Return to Engineer[/engineer]", border_style="engineer"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
elif fn == "list_nodes": obs = self.list_nodes_tool(**args)
|
||
elif fn == "run_commands": obs = self.run_commands_tool(**args, status=status)
|
||
elif fn == "get_node_info": obs = self.get_node_info_tool(**args)
|
||
elif fn == "manage_memory_tool": obs = self.manage_memory_tool(**args)
|
||
elif fn.startswith("mcp_"):
|
||
obs = run_ai_async(self.mcp_manager.call_tool(fn, args)).result(timeout=60)
|
||
elif fn in self.external_tool_handlers: obs = self.external_tool_handlers[fn](self, **args)
|
||
else: obs = f"Error: {fn} unknown."
|
||
|
||
if debug and fn not in ["delegate_to_engineer", "consult_architect", "escalate_to_architect", "return_to_engineer"]:
|
||
self._print_debug_observation(f"Observation: {fn}", obs, status=status)
|
||
|
||
# Ensure observation is a string and truncated for the LLM
|
||
obs_str = obs if isinstance(obs, str) else json.dumps(obs)
|
||
messages.append({"tool_call_id": tc.id, "role": "tool", "name": fn, "content": self._truncate(obs_str)})
|
||
# Inject pending user message AFTER all tool responses are added
|
||
if pending_user_message:
|
||
messages.append({"role": "user", "content": pending_user_message})
|
||
|
||
if iteration >= self.hard_limit_iterations:
|
||
self.console.print(f"[error]⛔ Agent reached hard limit ({self.hard_limit_iterations} steps). Forcing stop to prevent infinite loop.[/error]")
|
||
# Only inject user message if we're not in the middle of tool calls
|
||
last_msg = messages[-1] if messages else {}
|
||
if last_msg.get("role") != "assistant" or not last_msg.get("tool_calls"):
|
||
messages.append({"role": "user", "content": "Hard iteration limit reached. Please provide a summary of your findings so far."})
|
||
try:
|
||
safe_messages = self._sanitize_messages(messages)
|
||
response = completion(model=model, messages=safe_messages, tools=[], api_key=key)
|
||
resp_msg = response.choices[0].message
|
||
messages.append(resp_msg.model_dump(exclude_none=True))
|
||
except Exception as e:
|
||
if status:
|
||
status.update(f"[error]Error fetching summary: {e}[/error]")
|
||
printer.warning(f"Failed to fetch final summary from LLM: {e}")
|
||
except KeyboardInterrupt:
|
||
if status: status.update("[error]Interrupted! Closing pending tasks...")
|
||
last_msg = messages[-1]
|
||
if last_msg.get("tool_calls"):
|
||
for tc in last_msg["tool_calls"]:
|
||
messages.append({"tool_call_id": tc.get("id"), "role": "tool", "name": tc.get("function", {}).get("name"), "content": "Operation cancelled by user."})
|
||
|
||
# Use a fresh list for the summary call to avoid history corruption
|
||
summary_messages = list(messages)
|
||
summary_messages.append({"role": "user", "content": "USER INTERRUPTED. Briefly summarize what you were doing and stop."})
|
||
try:
|
||
safe_messages = self._sanitize_messages(summary_messages)
|
||
# Use tools=None to force a text summary during interruption
|
||
response = completion(model=model, messages=safe_messages, tools=None, api_key=key)
|
||
resp_msg = response.choices[0].message
|
||
messages.append(resp_msg.model_dump(exclude_none=True))
|
||
|
||
# IMPORTANT: Manually trigger callback for the summary so Web UI sees it
|
||
if chunk_callback and resp_msg.content:
|
||
chunk_callback(resp_msg.content)
|
||
except Exception:
|
||
error_msg = "Operation interrupted by user. Summary unavailable."
|
||
messages.append({"role": "assistant", "content": error_msg})
|
||
if chunk_callback:
|
||
chunk_callback(error_msg)
|
||
finally:
|
||
# Auto-save session
|
||
self.save_session(messages, model=model)
|
||
|
||
return {
|
||
"response": messages[-1].get("content"),
|
||
"chat_history": messages[1:],
|
||
"app_related": True,
|
||
"usage": usage,
|
||
"responder": current_brain, # "architect" or "engineer"
|
||
"streamed": streamed_response
|
||
}
|
||
|
||
@MethodHook
|
||
async def aask_copilot(self, terminal_buffer, user_question, node_info=None, chunk_callback=None):
|
||
import json
|
||
import re
|
||
from litellm import acompletion
|
||
import asyncio
|
||
import warnings
|
||
import aiohttp
|
||
|
||
# Suppress unawaited coroutine warnings from LiteLLM's internal streaming logic during sudden cancellation
|
||
warnings.filterwarnings("ignore", message="coroutine '.*async_streaming.*' was never awaited", category=RuntimeWarning)
|
||
|
||
node_info = node_info or {}
|
||
os_info = node_info.get("os", "unknown")
|
||
node_name = node_info.get("name", "unknown")
|
||
persona = node_info.get("persona", "engineer")
|
||
memories = node_info.get("memories", [])
|
||
|
||
vendor_reference = ""
|
||
if os_info and os_info != "unknown":
|
||
try:
|
||
os_filename = os_info.lower().replace(" ", "_")
|
||
ref_path = os.path.join(self.config.defaultdir, "ai_references", f"{os_filename}.md")
|
||
if os.path.exists(ref_path):
|
||
with open(ref_path, "r") as f:
|
||
vendor_reference = f.read().strip()
|
||
except Exception:
|
||
pass
|
||
|
||
if persona == "architect":
|
||
system_prompt = f"""Role: NETWORK ARCHITECT. You act as a senior strategic advisor during a live SSH session.
|
||
Rules:
|
||
1. Answer the user's question directly based on the Terminal Context.
|
||
2. Focus on the "why" and "how". Analyze topologies, design patterns, and validate configurations.
|
||
3. Do NOT provide commands to execute unless specifically requested. Instead, explain the consequences and best practices.
|
||
4. Keep your guide concise and authoritative.
|
||
5. You MUST output your response in the following strict format:
|
||
<guide>
|
||
Your brief tactical guide in markdown.
|
||
</guide>
|
||
<commands>
|
||
</commands>
|
||
<risk>
|
||
low
|
||
</risk>
|
||
6. Risk level is usually "low" for read-only/no commands.
|
||
|
||
Terminal Context:
|
||
{terminal_buffer}
|
||
|
||
Device OS: {os_info}
|
||
Node: {node_name}"""
|
||
else:
|
||
system_prompt = f"""Role: TERMINAL COPILOT. You assist a network engineer during a live SSH session.
|
||
Rules:
|
||
1. Answer the user's question directly based on the Terminal Context.
|
||
2. If the user asks you to analyze, parse, or extract data from the Terminal Context, DO IT directly in the <guide> section (you can use markdown tables or lists). Do NOT just give them a command to do it themselves.
|
||
3. If the user wants to execute an action, provide the required CLI commands inside a <commands> block, one command per line. If no commands are needed, leave it empty or omit the block.
|
||
4. ULTRA-CONCISE. Keep your guide to the point.
|
||
5. You MUST output your response in the following strict format:
|
||
<guide>
|
||
Your brief tactical guide in markdown. 3-4 sentences max.
|
||
</guide>
|
||
<commands>
|
||
command 1
|
||
command 2
|
||
</commands>
|
||
<risk>
|
||
low, high, or destructive
|
||
</risk>
|
||
6. Risk level: "low" for read-only/no commands, "high" for config changes, "destructive" for potentially dangerous ops.
|
||
|
||
Terminal Context:
|
||
{terminal_buffer}
|
||
|
||
Device OS: {os_info}
|
||
Node: {node_name}"""
|
||
|
||
if vendor_reference:
|
||
system_prompt += f"\n\nVendor Command Reference:\n{vendor_reference}"
|
||
|
||
if memories:
|
||
system_prompt += "\n\nSession Memory (Important Facts):\n"
|
||
for m in memories:
|
||
system_prompt += f"- {m}\n"
|
||
|
||
# Fetch MCP tools for the current OS
|
||
mcp_tools = []
|
||
try:
|
||
mcp_tools = await self.mcp_manager.get_tools_for_llm(os_filter=os_info)
|
||
except Exception:
|
||
pass
|
||
|
||
if mcp_tools:
|
||
system_prompt += f"\n\nAvailable MCP Tools: {', '.join([t['function']['name'] for t in mcp_tools])}"
|
||
system_prompt += "\nUse these tools to validate syntax or find exact commands if needed before providing the final guide."
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_question}
|
||
]
|
||
|
||
iteration = 0
|
||
max_iterations = 5 # Allow up to 5 iterations for tool usage
|
||
|
||
# Use models based on persona
|
||
current_model = self.architect_model if persona == "architect" else self.engineer_model
|
||
current_key = self.architect_key if persona == "architect" else self.engineer_key
|
||
|
||
try:
|
||
while iteration < max_iterations:
|
||
iteration += 1
|
||
response = await acompletion(
|
||
model=current_model,
|
||
messages=messages,
|
||
tools=mcp_tools if mcp_tools else None,
|
||
api_key=current_key,
|
||
stream=True
|
||
)
|
||
|
||
full_content = ""
|
||
streamed_guide = ""
|
||
tool_calls = []
|
||
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
|
||
# Accumulate tool calls
|
||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||
for tc in delta.tool_calls:
|
||
idx = tc.index
|
||
if idx >= len(tool_calls):
|
||
tool_calls.append({"id": tc.id, "type": "function", "function": {"name": tc.function.name or "", "arguments": tc.function.arguments or ""}})
|
||
else:
|
||
if tc.id: tool_calls[idx]["id"] = tc.id
|
||
if tc.function.name: tool_calls[idx]["function"]["name"] = tc.function.name
|
||
if tc.function.arguments: tool_calls[idx]["function"]["arguments"] += tc.function.arguments
|
||
|
||
if hasattr(delta, 'content') and delta.content:
|
||
full_content += delta.content
|
||
|
||
if chunk_callback and not tool_calls: # Only stream if not using tools
|
||
start_idx = full_content.find("<guide>")
|
||
if start_idx != -1:
|
||
after_start = full_content[start_idx + 7:]
|
||
end_idx = after_start.find("</guide>")
|
||
|
||
if end_idx != -1:
|
||
current_guide = after_start[:end_idx]
|
||
else:
|
||
current_guide = after_start
|
||
if current_guide.endswith("<"): current_guide = current_guide[:-1]
|
||
elif current_guide.endswith("</"): current_guide = current_guide[:-2]
|
||
elif current_guide.endswith("</g"): current_guide = current_guide[:-3]
|
||
elif current_guide.endswith("</gu"): current_guide = current_guide[:-4]
|
||
elif current_guide.endswith("</gui"): current_guide = current_guide[:-5]
|
||
elif current_guide.endswith("</guid"): current_guide = current_guide[:-6]
|
||
elif current_guide.endswith("</guide"): current_guide = current_guide[:-7]
|
||
|
||
new_text = current_guide[len(streamed_guide):]
|
||
if new_text:
|
||
chunk_callback(new_text)
|
||
streamed_guide += new_text
|
||
|
||
if not tool_calls:
|
||
break
|
||
|
||
# Execute tool calls
|
||
messages.append({"role": "assistant", "content": full_content or None, "tool_calls": tool_calls})
|
||
for tc in tool_calls:
|
||
fn = tc["function"]["name"]
|
||
args = json.loads(tc["function"]["arguments"])
|
||
|
||
if "mcp_" in fn:
|
||
try:
|
||
obs = await asyncio.wait_for(self.mcp_manager.call_tool(fn, args), timeout=30.0)
|
||
except Exception as e:
|
||
obs = f"Error calling MCP tool: {e}"
|
||
else:
|
||
obs = f"Error: Tool {fn} not allowed in Copilot."
|
||
|
||
messages.append({"tool_call_id": tc["id"], "role": "tool", "name": fn, "content": self._truncate(str(obs))})
|
||
|
||
# If we hit the limit and it was still using tools, force a final answer
|
||
if tool_calls and iteration >= max_iterations:
|
||
messages.append({"role": "user", "content": "Tool limit reached. Provide your final tactical guide now based on the findings."})
|
||
response = await acompletion(
|
||
model=self.engineer_model,
|
||
messages=messages,
|
||
tools=None,
|
||
api_key=self.engineer_key,
|
||
stream=True
|
||
)
|
||
|
||
full_content = ""
|
||
streamed_guide = ""
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
if hasattr(delta, 'content') and delta.content:
|
||
full_content += delta.content
|
||
if chunk_callback:
|
||
start_idx = full_content.find("<guide>")
|
||
if start_idx != -1:
|
||
after_start = full_content[start_idx + 7:]
|
||
end_idx = after_start.find("</guide>")
|
||
if end_idx != -1:
|
||
current_guide = after_start[:end_idx]
|
||
else:
|
||
current_guide = after_start
|
||
if current_guide.endswith("<"): current_guide = current_guide[:-1]
|
||
elif current_guide.endswith("</"): current_guide = current_guide[:-2]
|
||
elif current_guide.endswith("</g"): current_guide = current_guide[:-3]
|
||
elif current_guide.endswith("</gu"): current_guide = current_guide[:-4]
|
||
elif current_guide.endswith("</gui"): current_guide = current_guide[:-5]
|
||
elif current_guide.endswith("</guid"): current_guide = current_guide[:-6]
|
||
elif current_guide.endswith("</guide"): current_guide = current_guide[:-7]
|
||
new_text = current_guide[len(streamed_guide):]
|
||
if new_text:
|
||
chunk_callback(new_text)
|
||
streamed_guide += new_text
|
||
|
||
guide = ""
|
||
commands = []
|
||
risk_level = "low"
|
||
|
||
guide_match = re.search(r"<guide>(.*?)</guide>", full_content, re.DOTALL)
|
||
if guide_match:
|
||
guide = guide_match.group(1).strip()
|
||
|
||
cmd_match = re.search(r"<commands>(.*?)</commands>", full_content, re.DOTALL)
|
||
if cmd_match:
|
||
cmds_raw = cmd_match.group(1).strip()
|
||
if cmds_raw:
|
||
commands = [c.strip() for c in cmds_raw.split('\n') if c.strip()]
|
||
|
||
risk_match = re.search(r"<risk>(.*?)</risk>", full_content, re.DOTALL)
|
||
if risk_match:
|
||
risk_level = risk_match.group(1).strip().lower()
|
||
|
||
if not guide and full_content and not ("<guide>" in full_content):
|
||
guide = full_content.strip()
|
||
|
||
return {
|
||
"commands": commands,
|
||
"guide": guide,
|
||
"risk_level": risk_level,
|
||
"error": None
|
||
}
|
||
|
||
except asyncio.CancelledError:
|
||
# Client cancelled the request via gRPC or local interrupt
|
||
if 'response' in locals():
|
||
try:
|
||
if hasattr(response, 'aclose'):
|
||
# Fire and forget the close to avoid blocking the cancel
|
||
asyncio.create_task(response.aclose())
|
||
elif hasattr(response, 'close'):
|
||
response.close()
|
||
except Exception:
|
||
pass
|
||
return None
|
||
except Exception as e:
|
||
return {
|
||
"commands": [],
|
||
"guide": "",
|
||
"risk_level": "low",
|
||
"error": str(e)
|
||
}
|
||
|
||
@MethodHook
|
||
def confirm(self, user_input): return True</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Hybrid Multi-Agent System: Selective Escalation with Role Persistence.</p></div>
|
||
<h3>Class variables</h3>
|
||
<dl>
|
||
<dt id="connpy.ai.SAFE_COMMANDS"><code class="name">var <span class="ident">SAFE_COMMANDS</span></code></dt>
|
||
<dd>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
</dl>
|
||
<h3>Instance variables</h3>
|
||
<dl>
|
||
<dt id="connpy.ai.architect_system_prompt"><code class="name">prop <span class="ident">architect_system_prompt</span></code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@property
|
||
def architect_system_prompt(self):
|
||
"""Build architect system prompt with plugin extensions."""
|
||
if self.architect_prompt_extensions:
|
||
extensions = "\n".join(self.architect_prompt_extensions)
|
||
return self._architect_base_prompt + f"\n\nPlugin Capabilities:\n{extensions}"
|
||
return self._architect_base_prompt</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Build architect system prompt with plugin extensions.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.engineer_system_prompt"><code class="name">prop <span class="ident">engineer_system_prompt</span></code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@property
|
||
def engineer_system_prompt(self):
|
||
"""Build engineer system prompt with plugin extensions."""
|
||
if self.engineer_prompt_extensions:
|
||
extensions = "\n".join(self.engineer_prompt_extensions)
|
||
return self._engineer_base_prompt + f"\n\nPlugin Capabilities:\n{extensions}"
|
||
return self._engineer_base_prompt</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Build engineer system prompt with plugin extensions.</p></div>
|
||
</dd>
|
||
</dl>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="connpy.ai.aask_copilot"><code class="name flex">
|
||
<span>async def <span class="ident">aask_copilot</span></span>(<span>self, terminal_buffer, user_question, node_info=None, chunk_callback=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python"> @MethodHook
|
||
async def aask_copilot(self, terminal_buffer, user_question, node_info=None, chunk_callback=None):
|
||
import json
|
||
import re
|
||
from litellm import acompletion
|
||
import asyncio
|
||
import warnings
|
||
import aiohttp
|
||
|
||
# Suppress unawaited coroutine warnings from LiteLLM's internal streaming logic during sudden cancellation
|
||
warnings.filterwarnings("ignore", message="coroutine '.*async_streaming.*' was never awaited", category=RuntimeWarning)
|
||
|
||
node_info = node_info or {}
|
||
os_info = node_info.get("os", "unknown")
|
||
node_name = node_info.get("name", "unknown")
|
||
persona = node_info.get("persona", "engineer")
|
||
memories = node_info.get("memories", [])
|
||
|
||
vendor_reference = ""
|
||
if os_info and os_info != "unknown":
|
||
try:
|
||
os_filename = os_info.lower().replace(" ", "_")
|
||
ref_path = os.path.join(self.config.defaultdir, "ai_references", f"{os_filename}.md")
|
||
if os.path.exists(ref_path):
|
||
with open(ref_path, "r") as f:
|
||
vendor_reference = f.read().strip()
|
||
except Exception:
|
||
pass
|
||
|
||
if persona == "architect":
|
||
system_prompt = f"""Role: NETWORK ARCHITECT. You act as a senior strategic advisor during a live SSH session.
|
||
Rules:
|
||
1. Answer the user's question directly based on the Terminal Context.
|
||
2. Focus on the "why" and "how". Analyze topologies, design patterns, and validate configurations.
|
||
3. Do NOT provide commands to execute unless specifically requested. Instead, explain the consequences and best practices.
|
||
4. Keep your guide concise and authoritative.
|
||
5. You MUST output your response in the following strict format:
|
||
<guide>
|
||
Your brief tactical guide in markdown.
|
||
</guide>
|
||
<commands>
|
||
</commands>
|
||
<risk>
|
||
low
|
||
</risk>
|
||
6. Risk level is usually "low" for read-only/no commands.
|
||
|
||
Terminal Context:
|
||
{terminal_buffer}
|
||
|
||
Device OS: {os_info}
|
||
Node: {node_name}"""
|
||
else:
|
||
system_prompt = f"""Role: TERMINAL COPILOT. You assist a network engineer during a live SSH session.
|
||
Rules:
|
||
1. Answer the user's question directly based on the Terminal Context.
|
||
2. If the user asks you to analyze, parse, or extract data from the Terminal Context, DO IT directly in the <guide> section (you can use markdown tables or lists). Do NOT just give them a command to do it themselves.
|
||
3. If the user wants to execute an action, provide the required CLI commands inside a <commands> block, one command per line. If no commands are needed, leave it empty or omit the block.
|
||
4. ULTRA-CONCISE. Keep your guide to the point.
|
||
5. You MUST output your response in the following strict format:
|
||
<guide>
|
||
Your brief tactical guide in markdown. 3-4 sentences max.
|
||
</guide>
|
||
<commands>
|
||
command 1
|
||
command 2
|
||
</commands>
|
||
<risk>
|
||
low, high, or destructive
|
||
</risk>
|
||
6. Risk level: "low" for read-only/no commands, "high" for config changes, "destructive" for potentially dangerous ops.
|
||
|
||
Terminal Context:
|
||
{terminal_buffer}
|
||
|
||
Device OS: {os_info}
|
||
Node: {node_name}"""
|
||
|
||
if vendor_reference:
|
||
system_prompt += f"\n\nVendor Command Reference:\n{vendor_reference}"
|
||
|
||
if memories:
|
||
system_prompt += "\n\nSession Memory (Important Facts):\n"
|
||
for m in memories:
|
||
system_prompt += f"- {m}\n"
|
||
|
||
# Fetch MCP tools for the current OS
|
||
mcp_tools = []
|
||
try:
|
||
mcp_tools = await self.mcp_manager.get_tools_for_llm(os_filter=os_info)
|
||
except Exception:
|
||
pass
|
||
|
||
if mcp_tools:
|
||
system_prompt += f"\n\nAvailable MCP Tools: {', '.join([t['function']['name'] for t in mcp_tools])}"
|
||
system_prompt += "\nUse these tools to validate syntax or find exact commands if needed before providing the final guide."
|
||
|
||
messages = [
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": user_question}
|
||
]
|
||
|
||
iteration = 0
|
||
max_iterations = 5 # Allow up to 5 iterations for tool usage
|
||
|
||
# Use models based on persona
|
||
current_model = self.architect_model if persona == "architect" else self.engineer_model
|
||
current_key = self.architect_key if persona == "architect" else self.engineer_key
|
||
|
||
try:
|
||
while iteration < max_iterations:
|
||
iteration += 1
|
||
response = await acompletion(
|
||
model=current_model,
|
||
messages=messages,
|
||
tools=mcp_tools if mcp_tools else None,
|
||
api_key=current_key,
|
||
stream=True
|
||
)
|
||
|
||
full_content = ""
|
||
streamed_guide = ""
|
||
tool_calls = []
|
||
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
|
||
# Accumulate tool calls
|
||
if hasattr(delta, 'tool_calls') and delta.tool_calls:
|
||
for tc in delta.tool_calls:
|
||
idx = tc.index
|
||
if idx >= len(tool_calls):
|
||
tool_calls.append({"id": tc.id, "type": "function", "function": {"name": tc.function.name or "", "arguments": tc.function.arguments or ""}})
|
||
else:
|
||
if tc.id: tool_calls[idx]["id"] = tc.id
|
||
if tc.function.name: tool_calls[idx]["function"]["name"] = tc.function.name
|
||
if tc.function.arguments: tool_calls[idx]["function"]["arguments"] += tc.function.arguments
|
||
|
||
if hasattr(delta, 'content') and delta.content:
|
||
full_content += delta.content
|
||
|
||
if chunk_callback and not tool_calls: # Only stream if not using tools
|
||
start_idx = full_content.find("<guide>")
|
||
if start_idx != -1:
|
||
after_start = full_content[start_idx + 7:]
|
||
end_idx = after_start.find("</guide>")
|
||
|
||
if end_idx != -1:
|
||
current_guide = after_start[:end_idx]
|
||
else:
|
||
current_guide = after_start
|
||
if current_guide.endswith("<"): current_guide = current_guide[:-1]
|
||
elif current_guide.endswith("</"): current_guide = current_guide[:-2]
|
||
elif current_guide.endswith("</g"): current_guide = current_guide[:-3]
|
||
elif current_guide.endswith("</gu"): current_guide = current_guide[:-4]
|
||
elif current_guide.endswith("</gui"): current_guide = current_guide[:-5]
|
||
elif current_guide.endswith("</guid"): current_guide = current_guide[:-6]
|
||
elif current_guide.endswith("</guide"): current_guide = current_guide[:-7]
|
||
|
||
new_text = current_guide[len(streamed_guide):]
|
||
if new_text:
|
||
chunk_callback(new_text)
|
||
streamed_guide += new_text
|
||
|
||
if not tool_calls:
|
||
break
|
||
|
||
# Execute tool calls
|
||
messages.append({"role": "assistant", "content": full_content or None, "tool_calls": tool_calls})
|
||
for tc in tool_calls:
|
||
fn = tc["function"]["name"]
|
||
args = json.loads(tc["function"]["arguments"])
|
||
|
||
if "mcp_" in fn:
|
||
try:
|
||
obs = await asyncio.wait_for(self.mcp_manager.call_tool(fn, args), timeout=30.0)
|
||
except Exception as e:
|
||
obs = f"Error calling MCP tool: {e}"
|
||
else:
|
||
obs = f"Error: Tool {fn} not allowed in Copilot."
|
||
|
||
messages.append({"tool_call_id": tc["id"], "role": "tool", "name": fn, "content": self._truncate(str(obs))})
|
||
|
||
# If we hit the limit and it was still using tools, force a final answer
|
||
if tool_calls and iteration >= max_iterations:
|
||
messages.append({"role": "user", "content": "Tool limit reached. Provide your final tactical guide now based on the findings."})
|
||
response = await acompletion(
|
||
model=self.engineer_model,
|
||
messages=messages,
|
||
tools=None,
|
||
api_key=self.engineer_key,
|
||
stream=True
|
||
)
|
||
|
||
full_content = ""
|
||
streamed_guide = ""
|
||
async for chunk in response:
|
||
delta = chunk.choices[0].delta
|
||
if hasattr(delta, 'content') and delta.content:
|
||
full_content += delta.content
|
||
if chunk_callback:
|
||
start_idx = full_content.find("<guide>")
|
||
if start_idx != -1:
|
||
after_start = full_content[start_idx + 7:]
|
||
end_idx = after_start.find("</guide>")
|
||
if end_idx != -1:
|
||
current_guide = after_start[:end_idx]
|
||
else:
|
||
current_guide = after_start
|
||
if current_guide.endswith("<"): current_guide = current_guide[:-1]
|
||
elif current_guide.endswith("</"): current_guide = current_guide[:-2]
|
||
elif current_guide.endswith("</g"): current_guide = current_guide[:-3]
|
||
elif current_guide.endswith("</gu"): current_guide = current_guide[:-4]
|
||
elif current_guide.endswith("</gui"): current_guide = current_guide[:-5]
|
||
elif current_guide.endswith("</guid"): current_guide = current_guide[:-6]
|
||
elif current_guide.endswith("</guide"): current_guide = current_guide[:-7]
|
||
new_text = current_guide[len(streamed_guide):]
|
||
if new_text:
|
||
chunk_callback(new_text)
|
||
streamed_guide += new_text
|
||
|
||
guide = ""
|
||
commands = []
|
||
risk_level = "low"
|
||
|
||
guide_match = re.search(r"<guide>(.*?)</guide>", full_content, re.DOTALL)
|
||
if guide_match:
|
||
guide = guide_match.group(1).strip()
|
||
|
||
cmd_match = re.search(r"<commands>(.*?)</commands>", full_content, re.DOTALL)
|
||
if cmd_match:
|
||
cmds_raw = cmd_match.group(1).strip()
|
||
if cmds_raw:
|
||
commands = [c.strip() for c in cmds_raw.split('\n') if c.strip()]
|
||
|
||
risk_match = re.search(r"<risk>(.*?)</risk>", full_content, re.DOTALL)
|
||
if risk_match:
|
||
risk_level = risk_match.group(1).strip().lower()
|
||
|
||
if not guide and full_content and not ("<guide>" in full_content):
|
||
guide = full_content.strip()
|
||
|
||
return {
|
||
"commands": commands,
|
||
"guide": guide,
|
||
"risk_level": risk_level,
|
||
"error": None
|
||
}
|
||
|
||
except asyncio.CancelledError:
|
||
# Client cancelled the request via gRPC or local interrupt
|
||
if 'response' in locals():
|
||
try:
|
||
if hasattr(response, 'aclose'):
|
||
# Fire and forget the close to avoid blocking the cancel
|
||
asyncio.create_task(response.aclose())
|
||
elif hasattr(response, 'close'):
|
||
response.close()
|
||
except Exception:
|
||
pass
|
||
return None
|
||
except Exception as e:
|
||
return {
|
||
"commands": [],
|
||
"guide": "",
|
||
"risk_level": "low",
|
||
"error": str(e)
|
||
}</code></pre>
|
||
</details>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt id="connpy.ai.ask"><code class="name flex">
|
||
<span>def <span class="ident">ask</span></span>(<span>self,<br>user_input,<br>dryrun=False,<br>chat_history=None,<br>status=None,<br>debug=False,<br>stream=True,<br>session_id=None,<br>chunk_callback=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def ask(self, user_input, dryrun=False, chat_history=None, status=None, debug=False, stream=True, session_id=None, chunk_callback=None):
|
||
if not self.engineer_key:
|
||
raise ValueError("Engineer API key not configured. Use 'connpy config --engineer-api-key <key>' to set it.")
|
||
|
||
if chat_history is None: chat_history = []
|
||
|
||
# Load session if provided and history is empty
|
||
if session_id and not chat_history:
|
||
session_data = self.load_session_data(session_id)
|
||
if session_data:
|
||
chat_history = session_data.get("history", [])
|
||
# If we loaded history, the caller might need it back
|
||
# But typically ask() is called in a loop with an external history object
|
||
|
||
usage = {"input": 0, "output": 0, "total": 0}
|
||
|
||
# 1. Selector de Rol inicial (Sticky Brain)
|
||
explicit_architect = re.match(r'^(architect|arquitecto|@architect)[:\s]', user_input, re.I)
|
||
explicit_engineer = re.match(r'^(engineer|ingeniero|@engineer)[:\s]', user_input, re.I)
|
||
|
||
if explicit_architect:
|
||
current_brain = "architect"
|
||
elif explicit_engineer:
|
||
current_brain = "engineer"
|
||
else:
|
||
# Sticky Brain: Detectar si el Arquitecto estaba al mando en el historial reciente
|
||
is_architect_active = False
|
||
for msg in reversed(chat_history[-5:]):
|
||
tcs = msg.get('tool_calls') if isinstance(msg, dict) else getattr(msg, 'tool_calls', None)
|
||
if tcs:
|
||
for tc in tcs:
|
||
fn = tc.get('function', {}).get('name') if isinstance(tc, dict) else getattr(getattr(tc, 'function', None), 'name', '')
|
||
# Architect stays in control if delegating tasks or if Engineer escalated to them
|
||
# consult_architect is just Engineer asking for advice - Engineer keeps control
|
||
if fn in ['delegate_to_engineer', 'escalate_to_architect']:
|
||
is_architect_active = True; break
|
||
if is_architect_active: break
|
||
current_brain = "architect" if is_architect_active else "engineer"
|
||
|
||
# 2. Preparación de mensajes y limpieza
|
||
clean_input = re.sub(r'^(architect|arquitecto|engineer|ingeniero|@architect|@engineer)[:\s]+', '', user_input, flags=re.IGNORECASE).strip()
|
||
|
||
system_prompt = self.architect_system_prompt if current_brain == "architect" else self.engineer_system_prompt
|
||
tools = self._get_architect_tools() if current_brain == "architect" else self._get_engineer_tools()
|
||
model = self.architect_model if current_brain == "architect" else self.engineer_model
|
||
key = self.architect_key if current_brain == "architect" else self.engineer_key
|
||
|
||
# Estructura optimizada para Prompt Caching (Solo para Anthropic directo, Vertex tiene reglas distintas)
|
||
if "claude" in model.lower() and "vertex" not in model.lower():
|
||
messages = [{"role": "system", "content": [{"type": "text", "text": system_prompt, "cache_control": {"type": "ephemeral"}}]}]
|
||
else:
|
||
messages = [{"role": "system", "content": system_prompt}]
|
||
|
||
# Interleaving de historial
|
||
last_role = "system"
|
||
# Sanitize history if the current target model is not compatible with cache_control
|
||
history_to_process = chat_history[-self.max_history:]
|
||
if "claude" not in model.lower() or "vertex" in model.lower():
|
||
history_to_process = self._sanitize_messages(history_to_process)
|
||
|
||
for msg in history_to_process:
|
||
m = msg if isinstance(msg, dict) else msg.model_dump(exclude_none=True)
|
||
role = m.get('role')
|
||
if role == last_role and role == 'user':
|
||
messages[-1]['content'] += "\n" + (m.get('content') or "")
|
||
continue
|
||
if role == 'assistant' and m.get('tool_calls') and m.get('content') == "": m['content'] = None
|
||
messages.append(m)
|
||
last_role = role
|
||
|
||
if last_role == 'user': messages[-1]['content'] += "\n" + clean_input
|
||
else: messages.append({"role": "user", "content": clean_input})
|
||
|
||
# 3. Bucle de ejecución
|
||
iteration = 0
|
||
try:
|
||
# Set up remote interrupt callback if bridge is provided
|
||
if status and hasattr(status, "on_interrupt"):
|
||
status.on_interrupt = lambda: setattr(self, "interrupted", True)
|
||
|
||
while iteration < self.hard_limit_iterations:
|
||
iteration += 1
|
||
|
||
# Check for interruption
|
||
if self.interrupted:
|
||
raise KeyboardInterrupt
|
||
|
||
# Soft limit warning
|
||
if iteration == self.soft_limit_iterations and not soft_limit_warned:
|
||
self.console.print(f"[warning]⚠ Agent has performed {iteration} steps. This is taking longer than expected.[/warning]")
|
||
self.console.print(f"[warning] You can press Ctrl+C to interrupt and get a summary of progress.[/warning]")
|
||
soft_limit_warned = True
|
||
|
||
label = "[architect][bold]Architect[/bold][/architect]" if current_brain == "architect" else "[engineer][bold]Engineer[/bold][/engineer]"
|
||
if status:
|
||
# Notify responder identity ONLY for web/remote clients (StatusBridge has is_web)
|
||
if getattr(status, "is_web", False):
|
||
status.update(f"__RESPONDER__:{current_brain}")
|
||
status.update(f"{label} is thinking... (step {iteration})")
|
||
|
||
streamed_response = False
|
||
try:
|
||
safe_messages = self._sanitize_messages(messages)
|
||
if stream:
|
||
response, streamed_response = self._stream_completion(
|
||
model=model, messages=safe_messages, tools=tools, api_key=key,
|
||
status=status, label=label, debug=debug, num_retries=3,
|
||
chunk_callback=chunk_callback
|
||
)
|
||
else:
|
||
response = completion(model=model, messages=safe_messages, tools=tools, api_key=key, num_retries=3)
|
||
except Exception as e:
|
||
if current_brain == "architect":
|
||
if status: status.update("[unavailable]Architect unavailable! Falling back to Engineer...")
|
||
# Preserve context when falling back - use clean_input directly
|
||
current_brain = "engineer"
|
||
model = self.engineer_model
|
||
tools = self._get_engineer_tools()
|
||
key = self.engineer_key
|
||
# Rebuild messages with Engineer system prompt and original user request
|
||
messages = [{"role": "system", "content": self.engineer_system_prompt}]
|
||
# Add chat history if exists (excluding system prompt)
|
||
if chat_history:
|
||
for msg in chat_history[-self.max_history:]:
|
||
if msg.get('role') != 'system':
|
||
messages.append(msg)
|
||
# Add current user request
|
||
messages.append({"role": "user", "content": clean_input})
|
||
continue
|
||
else:
|
||
return {"response": f"Error: Both engines failed. {str(e)}", "chat_history": messages[1:], "usage": usage}
|
||
|
||
if hasattr(response, "usage") and response.usage:
|
||
usage["input"] += getattr(response.usage, "prompt_tokens", 0)
|
||
usage["output"] += getattr(response.usage, "completion_tokens", 0)
|
||
usage["total"] += getattr(response.usage, "total_tokens", 0)
|
||
|
||
resp_msg = response.choices[0].message
|
||
msg_dict = resp_msg.model_dump(exclude_none=True)
|
||
if msg_dict.get("tool_calls") and msg_dict.get("content") == "": msg_dict["content"] = None
|
||
messages.append(msg_dict)
|
||
|
||
if debug and resp_msg.content and not streamed_response:
|
||
# In CLI debug mode, only print intermediate reasoning if there are tool calls AND it wasn't already streamed.
|
||
# If there are no tool calls, this content is the final answer and will be printed by the caller.
|
||
if resp_msg.tool_calls:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Markdown(resp_msg.content), title=f"[{current_brain}][bold]{label} Reasoning[/bold][/{current_brain}]", border_style="architect" if current_brain == "architect" else "engineer"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
if not resp_msg.tool_calls: break
|
||
|
||
# Track if we need to inject a user message after all tool responses
|
||
pending_user_message = None
|
||
|
||
for tc in resp_msg.tool_calls:
|
||
fn, args = tc.function.name, json.loads(tc.function.arguments)
|
||
|
||
# Validate tool access based on current brain
|
||
if fn in ['delegate_to_engineer'] and current_brain != "architect":
|
||
obs = f"Error: Tool '{fn}' is only available to the Architect (Architect). You are the Engineer (Engineer). Use 'run_commands' directly to execute configuration."
|
||
messages.append({"tool_call_id": tc.id, "role": "tool", "name": fn, "content": obs})
|
||
continue
|
||
|
||
if status:
|
||
if fn == "delegate_to_engineer": status.update(f"[architect]Architect: [DELEGATING MISSION] {args.get('task','')[:40]}...")
|
||
elif fn == "manage_memory_tool": status.update(f"[architect]Architect: [UPDATING MEMORY]")
|
||
|
||
if debug:
|
||
self._print_debug_observation(f"Decision: {fn}", args, status=status)
|
||
|
||
if fn == "delegate_to_engineer":
|
||
obs, eng_usage = self._engineer_loop(args["task"], status=status, debug=debug, chat_history=messages[:-1])
|
||
usage["input"] += eng_usage["input"]; usage["output"] += eng_usage["output"]; usage["total"] += eng_usage["total"]
|
||
elif fn == "consult_architect":
|
||
if status: status.update("[architect]Engineer consulting Architect...")
|
||
try:
|
||
# Consultation only - Engineer stays in control
|
||
claude_resp = completion(
|
||
model=self.architect_model,
|
||
messages=[
|
||
{"role": "system", "content": self.architect_system_prompt},
|
||
{"role": "user", "content": f"The Engineer needs your strategic advice.\n\nTECHNICAL SUMMARY: {args['technical_summary']}\n\nQUESTION: {args['question']}\n\nProvide strategic guidance. The Engineer will continue handling the user."}
|
||
],
|
||
api_key=self.architect_key,
|
||
num_retries=3
|
||
)
|
||
obs = claude_resp.choices[0].message.content
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Markdown(obs), title="[architect]Architect Consultation[/architect]", border_style="architect"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
except Exception as e:
|
||
if status: status.update("[unavailable]Architect unavailable! Engineer continuing alone...")
|
||
obs = f"Architect unavailable ({str(e)}). Proceeding with your best technical judgment."
|
||
|
||
elif fn == "escalate_to_architect":
|
||
if status: status.update("[architect]Transferring control to Architect...")
|
||
# Full escalation - Architect takes over
|
||
current_brain = "architect"
|
||
model = self.architect_model
|
||
tools = self._get_architect_tools()
|
||
key = self.architect_key
|
||
messages[0] = {"role": "system", "content": self.architect_system_prompt}
|
||
# Prepare handover context to inject AFTER all tool responses
|
||
handover_msg = f"HANDOVER FROM EXECUTION ENGINE\n\nReason: {args['reason']}\n\nContext: {args['context']}\n\nYou are now in control of this conversation."
|
||
pending_user_message = handover_msg
|
||
obs = "Control transferred to Architect. Handover context will be provided."
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Text(handover_msg), title="[architect]Escalation to Architect[/architect]", border_style="architect"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
elif fn == "return_to_engineer":
|
||
if status: status.update("[engineer]Transferring control back to Engineer...")
|
||
# Architect returns control to Engineer
|
||
current_brain = "engineer"
|
||
model = self.engineer_model
|
||
tools = self._get_engineer_tools()
|
||
key = self.engineer_key
|
||
messages[0] = {"role": "system", "content": self.engineer_system_prompt}
|
||
# Prepare handover context to inject AFTER all tool responses
|
||
handover_msg = f"HANDOVER FROM ARCHITECT\n\nSummary: {args['summary']}\n\nYou are now back in control. Continue handling the user's requests."
|
||
pending_user_message = handover_msg
|
||
obs = "Control returned to Engineer. Handover summary will be provided."
|
||
if debug:
|
||
if status:
|
||
try: status.stop()
|
||
except: pass
|
||
self.console.print(Panel(Text(handover_msg), title="[engineer]Return to Engineer[/engineer]", border_style="engineer"))
|
||
if status:
|
||
try: status.start()
|
||
except: pass
|
||
|
||
elif fn == "list_nodes": obs = self.list_nodes_tool(**args)
|
||
elif fn == "run_commands": obs = self.run_commands_tool(**args, status=status)
|
||
elif fn == "get_node_info": obs = self.get_node_info_tool(**args)
|
||
elif fn == "manage_memory_tool": obs = self.manage_memory_tool(**args)
|
||
elif fn.startswith("mcp_"):
|
||
obs = run_ai_async(self.mcp_manager.call_tool(fn, args)).result(timeout=60)
|
||
elif fn in self.external_tool_handlers: obs = self.external_tool_handlers[fn](self, **args)
|
||
else: obs = f"Error: {fn} unknown."
|
||
|
||
if debug and fn not in ["delegate_to_engineer", "consult_architect", "escalate_to_architect", "return_to_engineer"]:
|
||
self._print_debug_observation(f"Observation: {fn}", obs, status=status)
|
||
|
||
# Ensure observation is a string and truncated for the LLM
|
||
obs_str = obs if isinstance(obs, str) else json.dumps(obs)
|
||
messages.append({"tool_call_id": tc.id, "role": "tool", "name": fn, "content": self._truncate(obs_str)})
|
||
# Inject pending user message AFTER all tool responses are added
|
||
if pending_user_message:
|
||
messages.append({"role": "user", "content": pending_user_message})
|
||
|
||
if iteration >= self.hard_limit_iterations:
|
||
self.console.print(f"[error]⛔ Agent reached hard limit ({self.hard_limit_iterations} steps). Forcing stop to prevent infinite loop.[/error]")
|
||
# Only inject user message if we're not in the middle of tool calls
|
||
last_msg = messages[-1] if messages else {}
|
||
if last_msg.get("role") != "assistant" or not last_msg.get("tool_calls"):
|
||
messages.append({"role": "user", "content": "Hard iteration limit reached. Please provide a summary of your findings so far."})
|
||
try:
|
||
safe_messages = self._sanitize_messages(messages)
|
||
response = completion(model=model, messages=safe_messages, tools=[], api_key=key)
|
||
resp_msg = response.choices[0].message
|
||
messages.append(resp_msg.model_dump(exclude_none=True))
|
||
except Exception as e:
|
||
if status:
|
||
status.update(f"[error]Error fetching summary: {e}[/error]")
|
||
printer.warning(f"Failed to fetch final summary from LLM: {e}")
|
||
except KeyboardInterrupt:
|
||
if status: status.update("[error]Interrupted! Closing pending tasks...")
|
||
last_msg = messages[-1]
|
||
if last_msg.get("tool_calls"):
|
||
for tc in last_msg["tool_calls"]:
|
||
messages.append({"tool_call_id": tc.get("id"), "role": "tool", "name": tc.get("function", {}).get("name"), "content": "Operation cancelled by user."})
|
||
|
||
# Use a fresh list for the summary call to avoid history corruption
|
||
summary_messages = list(messages)
|
||
summary_messages.append({"role": "user", "content": "USER INTERRUPTED. Briefly summarize what you were doing and stop."})
|
||
try:
|
||
safe_messages = self._sanitize_messages(summary_messages)
|
||
# Use tools=None to force a text summary during interruption
|
||
response = completion(model=model, messages=safe_messages, tools=None, api_key=key)
|
||
resp_msg = response.choices[0].message
|
||
messages.append(resp_msg.model_dump(exclude_none=True))
|
||
|
||
# IMPORTANT: Manually trigger callback for the summary so Web UI sees it
|
||
if chunk_callback and resp_msg.content:
|
||
chunk_callback(resp_msg.content)
|
||
except Exception:
|
||
error_msg = "Operation interrupted by user. Summary unavailable."
|
||
messages.append({"role": "assistant", "content": error_msg})
|
||
if chunk_callback:
|
||
chunk_callback(error_msg)
|
||
finally:
|
||
# Auto-save session
|
||
self.save_session(messages, model=model)
|
||
|
||
return {
|
||
"response": messages[-1].get("content"),
|
||
"chat_history": messages[1:],
|
||
"app_related": True,
|
||
"usage": usage,
|
||
"responder": current_brain, # "architect" or "engineer"
|
||
"streamed": streamed_response
|
||
}</code></pre>
|
||
</details>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt id="connpy.ai.confirm"><code class="name flex">
|
||
<span>def <span class="ident">confirm</span></span>(<span>self, user_input)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def confirm(self, user_input): return True</code></pre>
|
||
</details>
|
||
<div class="desc"></div>
|
||
</dd>
|
||
<dt id="connpy.ai.delete_session"><code class="name flex">
|
||
<span>def <span class="ident">delete_session</span></span>(<span>self, session_id)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def delete_session(self, session_id):
|
||
"""Deletes a session by ID."""
|
||
path = os.path.join(self.sessions_dir, f"{session_id}.json")
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
printer.success(f"Session {session_id} deleted.")
|
||
else:
|
||
printer.error(f"Session {session_id} not found.")</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Deletes a session by ID.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.get_last_session_id"><code class="name flex">
|
||
<span>def <span class="ident">get_last_session_id</span></span>(<span>self)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def get_last_session_id(self):
|
||
"""Returns the ID of the most recent session."""
|
||
sessions = self._get_sessions()
|
||
return sessions[0]["id"] if sessions else None</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Returns the ID of the most recent session.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.get_node_info_tool"><code class="name flex">
|
||
<span>def <span class="ident">get_node_info_tool</span></span>(<span>self, node_name)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def get_node_info_tool(self, node_name):
|
||
"""Get detailed metadata for a specific node. Passwords are masked."""
|
||
try:
|
||
d = self.config.getitem(node_name, extract=True)
|
||
if 'password' in d: d['password'] = '***'
|
||
return d
|
||
except Exception as e:
|
||
return f"Error getting node info: {str(e)}"</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get detailed metadata for a specific node. Passwords are masked.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.list_nodes_tool"><code class="name flex">
|
||
<span>def <span class="ident">list_nodes_tool</span></span>(<span>self, filter_pattern='.*')</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def list_nodes_tool(self, filter_pattern=".*"):
|
||
"""List nodes matching the filter pattern. Returns metadata for <=5 nodes, names only for more."""
|
||
try:
|
||
matched_names = self.config._getallnodes(filter_pattern)
|
||
if not matched_names: return "No nodes found."
|
||
if len(matched_names) <= 5:
|
||
matched_data = self.config.getitems(matched_names, extract=True)
|
||
res = {}
|
||
for name, data in matched_data.items():
|
||
os_tag = "unknown"
|
||
if isinstance(data, dict):
|
||
ts = data.get("tags")
|
||
if isinstance(ts, dict): os_tag = ts.get("os", "unknown")
|
||
res[name] = {"os": os_tag}
|
||
return res
|
||
return {"count": len(matched_names), "nodes": matched_names, "note": "Use 'get_node_info' for details."}
|
||
except Exception as e:
|
||
return f"Error listing nodes: {str(e)}"</code></pre>
|
||
</details>
|
||
<div class="desc"><p>List nodes matching the filter pattern. Returns metadata for <=5 nodes, names only for more.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.list_sessions"><code class="name flex">
|
||
<span>def <span class="ident">list_sessions</span></span>(<span>self)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def list_sessions(self):
|
||
"""Prints a list of sessions using printer.table."""
|
||
sessions = self._get_sessions()
|
||
if not sessions:
|
||
printer.info("No saved AI sessions found.")
|
||
return
|
||
|
||
columns = ["ID", "Title", "Created At", "Model"]
|
||
rows = [[s["id"], s["title"], s["created_at"], s["model"]] for s in sessions]
|
||
printer.table("AI Persisted Sessions", columns, rows)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Prints a list of sessions using printer.table.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.load_session_data"><code class="name flex">
|
||
<span>def <span class="ident">load_session_data</span></span>(<span>self, session_id)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def load_session_data(self, session_id):
|
||
"""Loads a session's raw data by ID."""
|
||
path = os.path.join(self.sessions_dir, f"{session_id}.json")
|
||
if os.path.exists(path):
|
||
try:
|
||
with open(path, "r") as f:
|
||
data = json.load(f)
|
||
self.session_id = session_id
|
||
self.session_path = path
|
||
return data
|
||
except Exception as e:
|
||
printer.error(f"Failed to load session {session_id}: {e}")
|
||
return None</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Loads a session's raw data by ID.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.manage_memory_tool"><code class="name flex">
|
||
<span>def <span class="ident">manage_memory_tool</span></span>(<span>self, content, action='append')</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def manage_memory_tool(self, content, action="append"):
|
||
"""Save or update long-term memory. Only use when user explicitly requests it."""
|
||
if not content or not content.strip():
|
||
return "Error: Cannot save empty content to memory."
|
||
|
||
try:
|
||
mode = "a" if action == "append" else "w"
|
||
os.makedirs(os.path.dirname(self.memory_path), exist_ok=True)
|
||
with open(self.memory_path, mode) as f:
|
||
timestamp = datetime.datetime.now().strftime('%Y-%m-%d %H:%M')
|
||
f.write(f"\n\n## {timestamp}\n{content.strip()}\n" if action == "append" else content)
|
||
|
||
# Reload memory after update
|
||
with open(self.memory_path, "r") as f:
|
||
self.long_term_memory = f.read()
|
||
|
||
return "Memory updated successfully."
|
||
except PermissionError as e:
|
||
return f"Error: Permission denied writing to memory file: {e}"
|
||
except Exception as e:
|
||
return f"Error updating memory: {str(e)}"</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Save or update long-term memory. Only use when user explicitly requests it.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.register_ai_tool"><code class="name flex">
|
||
<span>def <span class="ident">register_ai_tool</span></span>(<span>self,<br>tool_definition,<br>handler,<br>target='engineer',<br>engineer_prompt=None,<br>architect_prompt=None,<br>status_formatter=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def register_ai_tool(self, tool_definition, handler, target="engineer", engineer_prompt=None, architect_prompt=None, status_formatter=None):
|
||
"""Register an external tool for the AI system.
|
||
|
||
Args:
|
||
tool_definition (dict): OpenAI-compatible tool definition.
|
||
handler (callable): Function(ai_instance, **tool_args) -> str.
|
||
target (str): 'engineer', 'architect', or 'both'.
|
||
engineer_prompt (str): Extra text for engineer system prompt.
|
||
architect_prompt (str): Extra text for architect system prompt.
|
||
status_formatter (callable): Function(args_dict) -> status string.
|
||
"""
|
||
name = tool_definition["function"]["name"]
|
||
|
||
# Check if already registered to prevent duplicates
|
||
if target in ("engineer", "both"):
|
||
if not any(t["function"]["name"] == name for t in self.external_engineer_tools):
|
||
self.external_engineer_tools.append(tool_definition)
|
||
if target in ("architect", "both"):
|
||
if not any(t["function"]["name"] == name for t in self.external_architect_tools):
|
||
self.external_architect_tools.append(tool_definition)
|
||
|
||
self.external_tool_handlers[name] = handler
|
||
|
||
if engineer_prompt and engineer_prompt not in self.engineer_prompt_extensions:
|
||
self.engineer_prompt_extensions.append(engineer_prompt)
|
||
if architect_prompt and architect_prompt not in self.architect_prompt_extensions:
|
||
self.architect_prompt_extensions.append(architect_prompt)
|
||
if status_formatter:
|
||
self.tool_status_formatters[name] = status_formatter</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Register an external tool for the AI system.</p>
|
||
<h2 id="args">Args</h2>
|
||
<dl>
|
||
<dt><strong><code>tool_definition</code></strong> : <code>dict</code></dt>
|
||
<dd>OpenAI-compatible tool definition.</dd>
|
||
<dt><strong><code>handler</code></strong> : <code>callable</code></dt>
|
||
<dd>Function(ai_instance, **tool_args) -> str.</dd>
|
||
<dt><strong><code>target</code></strong> : <code>str</code></dt>
|
||
<dd>'engineer', 'architect', or 'both'.</dd>
|
||
<dt><strong><code>engineer_prompt</code></strong> : <code>str</code></dt>
|
||
<dd>Extra text for engineer system prompt.</dd>
|
||
<dt><strong><code>architect_prompt</code></strong> : <code>str</code></dt>
|
||
<dd>Extra text for architect system prompt.</dd>
|
||
<dt><strong><code>status_formatter</code></strong> : <code>callable</code></dt>
|
||
<dd>Function(args_dict) -> status string.</dd>
|
||
</dl></div>
|
||
</dd>
|
||
<dt id="connpy.ai.run_commands_tool"><code class="name flex">
|
||
<span>def <span class="ident">run_commands_tool</span></span>(<span>self, nodes_filter, commands, status=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def run_commands_tool(self, nodes_filter, commands, status=None):
|
||
"""Execute commands on nodes matching the filter. Native interactive confirmation for unsafe commands."""
|
||
# Handle if commands is a JSON string
|
||
if isinstance(commands, str):
|
||
try:
|
||
commands = json.loads(commands)
|
||
except ValueError:
|
||
commands = [c.strip() for c in commands.split('\n') if c.strip()]
|
||
|
||
# Expand multi-line commands within a list (in case the AI packs them)
|
||
if isinstance(commands, list):
|
||
expanded_commands = []
|
||
for cmd in commands:
|
||
expanded_commands.extend([c.strip() for c in str(cmd).split('\n') if c.strip()])
|
||
commands = expanded_commands
|
||
else:
|
||
commands = [str(commands)]
|
||
|
||
# Check command safety natively
|
||
if not self.trusted_session:
|
||
unsafe_commands = [cmd for cmd in commands if not self._is_safe_command(cmd)]
|
||
if unsafe_commands:
|
||
# Stop the spinner so prompt doesn't get messed up
|
||
if status: status.stop()
|
||
|
||
# Show ALL commands with unsafe ones highlighted
|
||
formatted_cmds = []
|
||
for cmd in commands:
|
||
if cmd in unsafe_commands:
|
||
formatted_cmds.append(f" • [warning]{cmd}[/warning]")
|
||
else:
|
||
formatted_cmds.append(f" • {cmd}")
|
||
|
||
panel_content = f"Target: {nodes_filter}\nCommands:\n" + "\n".join(formatted_cmds)
|
||
# Use print_important if available (for remote bridges) fallback to standard print
|
||
print_fn = getattr(self.console, "print_important", self.console.print)
|
||
print_fn(Panel(panel_content, title="[bold warning]⚠️ UNSAFE COMMANDS DETECTED[/bold warning]", border_style="warning"))
|
||
|
||
try:
|
||
user_resp = self.confirm_handler("[bold warning]Execute? (y: yes / n: no / a: allow all this session / <text>: feedback)[/bold warning]", default="n")
|
||
except KeyboardInterrupt:
|
||
if status: status.update("[ai_status]Engineer: Resuming...")
|
||
self.console.print("[fail]✗ Aborted by user (Ctrl+C).[/fail]")
|
||
raise
|
||
|
||
# Resume the spinner
|
||
if status: status.update("[ai_status]Engineer: Processing user response...")
|
||
|
||
user_resp_lower = user_resp.strip().lower()
|
||
if user_resp_lower in ['a', 'allow']:
|
||
self.trusted_session = True
|
||
self.console.print("[pass]✓ Trust Mode Enabled. All future commands in this session will execute without confirmation.[/pass]")
|
||
elif user_resp_lower in ['y', 'yes']:
|
||
self.console.print("[pass]✓ Executing...[/pass]")
|
||
elif user_resp_lower in ['n', 'no', '', 'cancel']:
|
||
self.console.print("[fail]✗ Execution rejected by user.[/fail]")
|
||
return "Error: User rejected execution."
|
||
else:
|
||
self.console.print(f"[user_prompt]User feedback: [/user_prompt]{user_resp}")
|
||
return f"User requested changes: {user_resp}. Please adjust the commands based on this feedback and try again."
|
||
|
||
try:
|
||
matched_names = self.config._getallnodes(nodes_filter)
|
||
if not matched_names: return "No nodes found matching filter."
|
||
thisnodes_dict = self.config.getitems(matched_names, extract=True)
|
||
result = nodes(thisnodes_dict, config=self.config).run(commands)
|
||
return result
|
||
except Exception as e:
|
||
return f"Error executing commands: {str(e)}"</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Execute commands on nodes matching the filter. Native interactive confirmation for unsafe commands.</p></div>
|
||
</dd>
|
||
<dt id="connpy.ai.save_session"><code class="name flex">
|
||
<span>def <span class="ident">save_session</span></span>(<span>self, history, title=None, model=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">def save_session(self, history, title=None, model=None):
|
||
"""Saves current history to the session file."""
|
||
if not self.session_id:
|
||
# Generate ID from first user query if available
|
||
first_user_msg = next((m["content"] for m in history if m["role"] == "user"), "new-session")
|
||
self.session_id = self._generate_session_id(first_user_msg)
|
||
self.session_path = os.path.join(self.sessions_dir, f"{self.session_id}.json")
|
||
|
||
# If it's a new file, we might want to set a better title
|
||
if not os.path.exists(self.session_path) and not title:
|
||
raw_title = next((m["content"] for m in history if m["role"] == "user"), "New Session")
|
||
# Clean title: remove newlines, multiple spaces
|
||
clean_title = " ".join(raw_title.split())
|
||
if len(clean_title) > 40:
|
||
title = clean_title[:37].strip() + "..."
|
||
else:
|
||
title = clean_title
|
||
|
||
try:
|
||
# Read existing metadata if it exists
|
||
metadata = {}
|
||
if os.path.exists(self.session_path):
|
||
with open(self.session_path, "r") as f:
|
||
metadata = json.load(f)
|
||
|
||
metadata.update({
|
||
"id": self.session_id,
|
||
"title": title or metadata.get("title", "New Session"),
|
||
"created_at": metadata.get("created_at", datetime.datetime.now().isoformat()),
|
||
"updated_at": datetime.datetime.now().isoformat(),
|
||
"model": model or metadata.get("model", self.engineer_model),
|
||
"history": history
|
||
})
|
||
|
||
with open(self.session_path, "w") as f:
|
||
json.dump(metadata, f, indent=4)
|
||
except Exception as e:
|
||
printer.error(f"Failed to save session: {e}")
|
||
|
||
except Exception as e:
|
||
printer.error(f"Failed to save session: {e}")</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Saves current history to the session file.</p></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
<dt id="connpy.configfile"><code class="flex name class">
|
||
<span>class <span class="ident">configfile</span></span>
|
||
<span>(</span><span>conf=None, key=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@ClassHook
|
||
class configfile:
|
||
''' This class generates a configfile object. Containts a dictionary storing, config, nodes and profiles, normaly used by connection manager.
|
||
|
||
### Attributes:
|
||
|
||
- file (str): Path/file to config file.
|
||
|
||
- key (str): Path/file to RSA key file.
|
||
|
||
- config (dict): Dictionary containing information of connection
|
||
manager configuration.
|
||
|
||
- connections (dict): Dictionary containing all the nodes added to
|
||
connection manager.
|
||
|
||
- profiles (dict): Dictionary containing all the profiles added to
|
||
connection manager.
|
||
|
||
- privatekey (obj): Object containing the private key to encrypt
|
||
passwords.
|
||
|
||
- publickey (obj): Object containing the public key to decrypt
|
||
passwords.
|
||
'''
|
||
|
||
def __init__(self, conf = None, key = None):
|
||
'''
|
||
|
||
### Optional Parameters:
|
||
|
||
- conf (str): Path/file to config file. If left empty default
|
||
path is ~/.config/conn/config.yaml
|
||
|
||
- key (str): Path/file to RSA key file. If left empty default
|
||
path is ~/.config/conn/.osk
|
||
|
||
'''
|
||
home = os.path.expanduser("~")
|
||
defaultdir = home + '/.config/conn'
|
||
|
||
if conf is None:
|
||
# Standard path: use ~/.config/conn and respect .folder redirection
|
||
self.anchor_path = defaultdir
|
||
self.defaultdir = defaultdir
|
||
Path(defaultdir).mkdir(parents=True, exist_ok=True)
|
||
|
||
pathfile = defaultdir + '/.folder'
|
||
try:
|
||
with open(pathfile, "r") as f:
|
||
configdir = f.read().strip()
|
||
except (FileNotFoundError, IOError):
|
||
with open(pathfile, "w") as f:
|
||
f.write(str(defaultdir))
|
||
configdir = defaultdir
|
||
|
||
self.defaultdir = configdir
|
||
self.file = configdir + '/config.yaml'
|
||
self.key = key or (configdir + '/.osk')
|
||
|
||
# Ensure redirected directories exist
|
||
Path(configdir).mkdir(parents=True, exist_ok=True)
|
||
Path(f"{configdir}/plugins").mkdir(parents=True, exist_ok=True)
|
||
|
||
# Backwards compatibility: Migrate from JSON to YAML only for default path
|
||
legacy_json = configdir + '/config.json'
|
||
legacy_noext = configdir + '/config'
|
||
legacy_file = None
|
||
if os.path.exists(legacy_json): legacy_file = legacy_json
|
||
elif os.path.exists(legacy_noext): legacy_file = legacy_noext
|
||
|
||
if not os.path.exists(self.file) and legacy_file:
|
||
try:
|
||
with open(legacy_file, 'r') as f:
|
||
old_data = json.load(f)
|
||
if not self._validate_config(old_data):
|
||
printer.warning(f"Legacy config {legacy_file} has invalid structure, skipping migration.")
|
||
else:
|
||
with open(self.file, 'w') as f:
|
||
yaml.dump(old_data, f, Dumper=NoAliasDumper, default_flow_style=False, sort_keys=False)
|
||
# Verify the written YAML can be read back correctly
|
||
with open(self.file, 'r') as f:
|
||
verify = yaml.safe_load(f)
|
||
if not self._validate_config(verify):
|
||
os.remove(self.file)
|
||
printer.warning("YAML verification failed after migration, keeping legacy config.")
|
||
else:
|
||
# Note: cachefile is derived later, we use temp one for migration sync
|
||
temp_cache = configdir + '/.config.cache.json'
|
||
with open(temp_cache, 'w') as f:
|
||
json.dump(old_data, f)
|
||
shutil.move(legacy_file, legacy_file + ".backup")
|
||
printer.success(f"Migrated legacy config ({len(old_data.get('connections',{}))} folders/nodes) into YAML and Cache successfully!")
|
||
except Exception as e:
|
||
if os.path.exists(self.file):
|
||
try: os.remove(self.file)
|
||
except OSError: pass
|
||
printer.warning(f"Failed to migrate legacy config: {e}")
|
||
else:
|
||
# Custom path (common in tests): isolate everything to the conf parent directory
|
||
self.file = os.path.abspath(conf)
|
||
configdir = os.path.dirname(self.file)
|
||
self.anchor_path = configdir
|
||
self.defaultdir = configdir
|
||
self.key = os.path.abspath(key) if key else (configdir + '/.osk')
|
||
|
||
# Sidecar files always live next to the config file (or in the redirected configdir)
|
||
self.cachefile = configdir + '/.config.cache.json'
|
||
self.fzf_cachefile = configdir + '/.fzf_nodes_cache.txt'
|
||
self.folders_cachefile = configdir + '/.folders_cache.txt'
|
||
self.profiles_cachefile = configdir + '/.profiles_cache.txt'
|
||
|
||
if os.path.exists(self.file):
|
||
config = self._loadconfig(self.file)
|
||
else:
|
||
config = self._createconfig(self.file)
|
||
|
||
self.config = config["config"]
|
||
self.connections = config["connections"]
|
||
self.profiles = config["profiles"]
|
||
|
||
if not os.path.exists(self.key):
|
||
self._createkey(self.key)
|
||
with open(self.key) as f:
|
||
self.privatekey = RSA.import_key(f.read())
|
||
self.publickey = self.privatekey.publickey()
|
||
|
||
# Self-heal text caches if they are missing
|
||
if not os.path.exists(self.fzf_cachefile) or not os.path.exists(self.folders_cachefile) or not os.path.exists(self.profiles_cachefile):
|
||
self._generate_nodes_cache()
|
||
|
||
|
||
def _validate_config(self, data):
|
||
"""Verify config data has the required structure."""
|
||
if not isinstance(data, dict):
|
||
return False
|
||
required = {"config", "connections", "profiles"}
|
||
return required.issubset(data.keys())
|
||
|
||
def _loadconfig(self, conf):
|
||
#Loads config file using dual cache
|
||
cache_exists = os.path.exists(self.cachefile)
|
||
yaml_time = os.path.getmtime(conf) if os.path.exists(conf) else 0
|
||
cache_time = os.path.getmtime(self.cachefile) if cache_exists else 0
|
||
|
||
if not cache_exists or yaml_time > cache_time:
|
||
with open(conf, 'r') as f:
|
||
data = yaml.safe_load(f)
|
||
if not self._validate_config(data):
|
||
# YAML is broken, try to recover from cache
|
||
if cache_exists:
|
||
printer.warning("Config file appears corrupt, recovering from cache...")
|
||
with open(self.cachefile, 'r') as f:
|
||
data = json.load(f)
|
||
if self._validate_config(data):
|
||
# Re-write the YAML from good cache
|
||
with open(conf, 'w') as f:
|
||
yaml.dump(data, f, Dumper=NoAliasDumper, default_flow_style=False, sort_keys=False)
|
||
return data
|
||
# Both broken or no cache - create fresh
|
||
printer.error("Config file is corrupt and no valid cache exists. Creating default config.")
|
||
return self._createconfig(conf)
|
||
try:
|
||
with open(self.cachefile, 'w') as f:
|
||
json.dump(data, f)
|
||
except Exception:
|
||
pass
|
||
return data
|
||
else:
|
||
with open(self.cachefile, 'r') as f:
|
||
data = json.load(f)
|
||
if not self._validate_config(data):
|
||
# Cache broken, try yaml
|
||
with open(conf, 'r') as f:
|
||
data = yaml.safe_load(f)
|
||
if self._validate_config(data):
|
||
return data
|
||
# Both broken
|
||
printer.error("Both config and cache are corrupt. Creating default config.")
|
||
return self._createconfig(conf)
|
||
return data
|
||
|
||
def _createconfig(self, conf):
|
||
#Create config file (always writes defaults, safe for recovery)
|
||
defaultconfig = {'config': {'case': False, 'idletime': 30, 'fzf': False}, 'connections': {}, 'profiles': { "default": { "host":"", "protocol":"ssh", "port":"", "user":"", "password":"", "options":"", "logs":"", "tags": "", "jumphost":""}}}
|
||
with open(conf, "w") as f:
|
||
yaml.dump(defaultconfig, f, Dumper=NoAliasDumper, default_flow_style=False, sort_keys=False)
|
||
os.chmod(conf, 0o600)
|
||
try:
|
||
with open(self.cachefile, 'w') as f:
|
||
json.dump(defaultconfig, f)
|
||
except Exception:
|
||
pass
|
||
return defaultconfig
|
||
|
||
@MethodHook
|
||
def _saveconfig(self, conf):
|
||
#Save config file atomically to prevent corruption
|
||
newconfig = {"config":{}, "connections": {}, "profiles": {}}
|
||
newconfig["config"] = self.config
|
||
newconfig["connections"] = self.connections
|
||
newconfig["profiles"] = self.profiles
|
||
tmpfile = conf + '.tmp'
|
||
try:
|
||
with open(tmpfile, "w") as f:
|
||
yaml.dump(newconfig, f, Dumper=NoAliasDumper, default_flow_style=False, sort_keys=False)
|
||
# Atomic replace: only overwrite original if write succeeded
|
||
shutil.move(tmpfile, conf)
|
||
with open(self.cachefile, "w") as f:
|
||
json.dump(newconfig, f)
|
||
self._generate_nodes_cache()
|
||
except (IOError, OSError) as e:
|
||
printer.error(f"Failed to save config: {e}")
|
||
# Clean up temp file if it exists
|
||
if os.path.exists(tmpfile):
|
||
try:
|
||
os.remove(tmpfile)
|
||
except OSError:
|
||
pass
|
||
return 1
|
||
return 0
|
||
|
||
def _generate_nodes_cache(self, nodes=None, folders=None, profiles=None):
|
||
try:
|
||
if nodes is None:
|
||
nodes = self._getallnodes()
|
||
if folders is None:
|
||
folders = self._getallfolders()
|
||
if profiles is None:
|
||
profiles = list(self.profiles.keys())
|
||
|
||
with open(self.fzf_cachefile, "w") as f:
|
||
f.write("\n".join(nodes))
|
||
with open(self.folders_cachefile, "w") as f:
|
||
f.write("\n".join(folders))
|
||
with open(self.profiles_cachefile, "w") as f:
|
||
f.write("\n".join(profiles))
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _createkey(self, keyfile):
|
||
#Create key file
|
||
key = RSA.generate(2048)
|
||
with open(keyfile,'wb') as f:
|
||
f.write(key.export_key('PEM'))
|
||
f.close()
|
||
os.chmod(keyfile, 0o600)
|
||
return key
|
||
|
||
@MethodHook
|
||
def _explode_unique(self, unique):
|
||
#Divide unique name into folder, subfolder and id
|
||
uniques = unique.split("@")
|
||
if not unique.startswith("@"):
|
||
result = {"id": uniques[0]}
|
||
else:
|
||
result = {}
|
||
if len(uniques) == 2:
|
||
result["folder"] = uniques[1]
|
||
if result["folder"] == "":
|
||
return False
|
||
elif len(uniques) == 3:
|
||
result["folder"] = uniques[2]
|
||
result["subfolder"] = uniques[1]
|
||
if result["folder"] == "" or result["subfolder"] == "":
|
||
return False
|
||
elif len(uniques) > 3:
|
||
return False
|
||
return result
|
||
|
||
@MethodHook
|
||
def getitem(self, unique, keys = None, extract = False):
|
||
'''
|
||
Get an node or a group of nodes from configfile which can be passed to node/nodes class
|
||
|
||
### Parameters:
|
||
|
||
- unique (str): Unique name of the node or folder in config using
|
||
connection manager style: node[@subfolder][@folder]
|
||
or [@subfolder]@folder
|
||
|
||
### Optional Parameters:
|
||
|
||
- keys (list): In case you pass a folder as unique, you can filter
|
||
nodes inside the folder passing a list.
|
||
- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
|
||
'''
|
||
uniques = self._explode_unique(unique)
|
||
if unique.startswith("@"):
|
||
if uniques.keys() >= {"folder", "subfolder"}:
|
||
folder = self.connections[uniques["folder"]][uniques["subfolder"]]
|
||
else:
|
||
folder = self.connections[uniques["folder"]]
|
||
newfolder = deepcopy(folder)
|
||
newfolder.pop("type")
|
||
for node_name in folder.keys():
|
||
if node_name == "type":
|
||
continue
|
||
if "type" in newfolder[node_name].keys():
|
||
if newfolder[node_name]["type"] == "subfolder":
|
||
newfolder.pop(node_name)
|
||
else:
|
||
newfolder[node_name].pop("type")
|
||
|
||
if keys != None:
|
||
newfolder = dict((k, newfolder[k]) for k in keys)
|
||
|
||
if extract:
|
||
for node_name, node_keys in newfolder.items():
|
||
for key, value in node_keys.items():
|
||
profile = re.search("^@(.*)", str(value))
|
||
if profile:
|
||
try:
|
||
newfolder[node_name][key] = self.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
newfolder[node_name][key] = ""
|
||
elif value == '' and key == "protocol":
|
||
try:
|
||
newfolder[node_name][key] = self.profiles["default"][key]
|
||
except KeyError:
|
||
newfolder[node_name][key] = "ssh"
|
||
|
||
newfolder = {"{}{}".format(k,unique):v for k,v in newfolder.items()}
|
||
return newfolder
|
||
else:
|
||
if uniques.keys() >= {"folder", "subfolder"}:
|
||
node = self.connections[uniques["folder"]][uniques["subfolder"]][uniques["id"]]
|
||
elif "folder" in uniques.keys():
|
||
node = self.connections[uniques["folder"]][uniques["id"]]
|
||
else:
|
||
node = self.connections[uniques["id"]]
|
||
newnode = deepcopy(node)
|
||
newnode.pop("type")
|
||
|
||
if extract:
|
||
for key, value in newnode.items():
|
||
profile = re.search("^@(.*)", str(value))
|
||
if profile:
|
||
try:
|
||
newnode[key] = self.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
newnode[key] = ""
|
||
elif value == '' and key == "protocol":
|
||
try:
|
||
newnode[key] = self.profiles["default"][key]
|
||
except KeyError:
|
||
newnode[key] = "ssh"
|
||
return newnode
|
||
|
||
@MethodHook
|
||
def getitems(self, uniques, extract = False):
|
||
'''
|
||
Get a group of nodes from configfile which can be passed to node/nodes class
|
||
|
||
### Parameters:
|
||
|
||
- uniques (str/list): String name that will match hostnames
|
||
from the connection manager. It can be a
|
||
list of strings.
|
||
|
||
### Optional Parameters:
|
||
|
||
- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
|
||
'''
|
||
nodes = {}
|
||
if isinstance(uniques, str):
|
||
uniques = [uniques]
|
||
for i in uniques:
|
||
if i.startswith("@"):
|
||
if not self.config["case"]:
|
||
i = i.lower()
|
||
this = self.getitem(i, extract = extract)
|
||
nodes.update(this)
|
||
else:
|
||
if not self.config["case"]:
|
||
i = i.lower()
|
||
this = self.getitem(i, extract = extract)
|
||
nodes[i] = this
|
||
return nodes
|
||
|
||
|
||
@MethodHook
|
||
def _connections_add(self,*, id, host, folder='', subfolder='', options='', logs='', password='', port='', protocol='', user='', tags='', jumphost='', type = "connection" ):
|
||
#Add connection from config
|
||
if folder == '':
|
||
self.connections[id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "tags": tags,"jumphost": jumphost,"type": type}
|
||
elif folder != '' and subfolder == '':
|
||
self.connections[folder][id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "tags": tags, "jumphost": jumphost, "type": type}
|
||
elif folder != '' and subfolder != '':
|
||
self.connections[folder][subfolder][id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "tags": tags, "jumphost": jumphost, "type": type}
|
||
|
||
|
||
@MethodHook
|
||
def _connections_del(self,*, id, folder='', subfolder=''):
|
||
#Delete connection from config
|
||
if folder == '':
|
||
del self.connections[id]
|
||
elif folder != '' and subfolder == '':
|
||
del self.connections[folder][id]
|
||
elif folder != '' and subfolder != '':
|
||
del self.connections[folder][subfolder][id]
|
||
|
||
@MethodHook
|
||
def _folder_add(self,*, folder, subfolder = ''):
|
||
#Add Folder from config
|
||
if subfolder == '':
|
||
if folder not in self.connections:
|
||
self.connections[folder] = {"type": "folder"}
|
||
else:
|
||
if subfolder not in self.connections[folder]:
|
||
self.connections[folder][subfolder] = {"type": "subfolder"}
|
||
|
||
@MethodHook
|
||
def _folder_del(self,*, folder, subfolder=''):
|
||
#Delete folder from config
|
||
if subfolder == '':
|
||
del self.connections[folder]
|
||
else:
|
||
del self.connections[folder][subfolder]
|
||
|
||
|
||
@MethodHook
|
||
def _profiles_add(self,*, id, host = '', options='', logs='', password='', port='', protocol='', user='', tags='', jumphost='' ):
|
||
#Add profile from config
|
||
self.profiles[id] = {"host": host, "options": options, "logs": logs, "password": password, "port": port, "protocol": protocol, "user": user, "tags": tags, "jumphost": jumphost}
|
||
|
||
|
||
@MethodHook
|
||
def _profiles_del(self,*, id ):
|
||
#Delete profile from config
|
||
del self.profiles[id]
|
||
|
||
@MethodHook
|
||
def _getallnodes(self, filter = None):
|
||
#get all nodes on configfile
|
||
nodes = []
|
||
layer1 = [k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "connection"]
|
||
folders = [k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "folder"]
|
||
nodes.extend(layer1)
|
||
for f in folders:
|
||
layer2 = [k + "@" + f for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "connection"]
|
||
nodes.extend(layer2)
|
||
subfolders = [k for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "subfolder"]
|
||
for s in subfolders:
|
||
layer3 = [k + "@" + s + "@" + f for k,v in self.connections[f][s].items() if isinstance(v, dict) and v.get("type") == "connection"]
|
||
nodes.extend(layer3)
|
||
if filter:
|
||
flat_filter = []
|
||
if isinstance(filter, str):
|
||
flat_filter = [filter]
|
||
elif isinstance(filter, list):
|
||
for item in filter:
|
||
if isinstance(item, str):
|
||
flat_filter.append(item)
|
||
else:
|
||
printer.error("Filter must be a string or a list of strings")
|
||
sys.exit(1)
|
||
nodes = [item for item in nodes if any(re.search(pattern, item) for pattern in flat_filter)]
|
||
return nodes
|
||
|
||
@MethodHook
|
||
def _getallnodesfull(self, filter = None, extract = True):
|
||
#get all nodes on configfile with all their attributes.
|
||
nodes = {}
|
||
layer1 = {k:v for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "connection"}
|
||
folders = [k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "folder"]
|
||
nodes.update(layer1)
|
||
for f in folders:
|
||
layer2 = {k + "@" + f:v for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "connection"}
|
||
nodes.update(layer2)
|
||
subfolders = [k for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "subfolder"]
|
||
for s in subfolders:
|
||
layer3 = {k + "@" + s + "@" + f:v for k,v in self.connections[f][s].items() if isinstance(v, dict) and v.get("type") == "connection"}
|
||
nodes.update(layer3)
|
||
if filter:
|
||
flat_filter = []
|
||
if isinstance(filter, str):
|
||
flat_filter = [filter]
|
||
elif isinstance(filter, list):
|
||
for item in filter:
|
||
if isinstance(item, str):
|
||
flat_filter.append(item)
|
||
else:
|
||
printer.error("Filter must be a string or a list of strings")
|
||
sys.exit(1)
|
||
flat_filter = ["^(?!.*@).+$" if item == "@" else item for item in flat_filter]
|
||
nodes = {k: v for k, v in nodes.items() if any(re.search(pattern, k) for pattern in flat_filter)}
|
||
if extract:
|
||
for node, keys in nodes.items():
|
||
for key, value in keys.items():
|
||
profile = re.search("^@(.*)", str(value))
|
||
if profile:
|
||
try:
|
||
nodes[node][key] = self.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
nodes[node][key] = ""
|
||
elif value == '' and key == "protocol":
|
||
try:
|
||
nodes[node][key] = self.profiles["default"][key]
|
||
except KeyError:
|
||
nodes[node][key] = "ssh"
|
||
return nodes
|
||
|
||
|
||
@MethodHook
|
||
def _getallfolders(self):
|
||
#get all folders on configfile
|
||
folders = ["@" + k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "folder"]
|
||
subfolders = []
|
||
for f in folders:
|
||
s = ["@" + k + f for k,v in self.connections[f[1:]].items() if isinstance(v, dict) and v.get("type") == "subfolder"]
|
||
subfolders.extend(s)
|
||
folders.extend(subfolders)
|
||
return folders
|
||
|
||
@MethodHook
|
||
def _profileused(self, profile):
|
||
#Return all the nodes that uses this profile.
|
||
nodes = []
|
||
layer1 = [k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "connection" and ("@" + profile in v.values() or ( isinstance(v.get("password"),list) and "@" + profile in v.get("password")))]
|
||
folders = [k for k,v in self.connections.items() if isinstance(v, dict) and v.get("type") == "folder"]
|
||
nodes.extend(layer1)
|
||
for f in folders:
|
||
layer2 = [k + "@" + f for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "connection" and ("@" + profile in v.values() or ( isinstance(v.get("password"),list) and "@" + profile in v.get("password")))]
|
||
nodes.extend(layer2)
|
||
subfolders = [k for k,v in self.connections[f].items() if isinstance(v, dict) and v.get("type") == "subfolder"]
|
||
for s in subfolders:
|
||
layer3 = [k + "@" + s + "@" + f for k,v in self.connections[f][s].items() if isinstance(v, dict) and v.get("type") == "connection" and ("@" + profile in v.values() or ( isinstance(v.get("password"),list) and "@" + profile in v.get("password")))]
|
||
nodes.extend(layer3)
|
||
return nodes
|
||
|
||
@MethodHook
|
||
def encrypt(self, password, keyfile=None):
|
||
'''
|
||
Encrypts password using RSA keyfile
|
||
|
||
### Parameters:
|
||
|
||
- password (str): Plaintext password to encrypt.
|
||
|
||
### Optional Parameters:
|
||
|
||
- keyfile (str): Path/file to keyfile. Default is config keyfile.
|
||
|
||
|
||
### Returns:
|
||
|
||
str: Encrypted password.
|
||
|
||
'''
|
||
if keyfile is None:
|
||
keyfile = self.key
|
||
with open(keyfile) as f:
|
||
key = RSA.import_key(f.read())
|
||
f.close()
|
||
publickey = key.publickey()
|
||
encryptor = PKCS1_OAEP.new(publickey)
|
||
password = encryptor.encrypt(password.encode("utf-8"))
|
||
return str(password)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>This class generates a configfile object. Containts a dictionary storing, config, nodes and profiles, normaly used by connection manager.</p>
|
||
<h3 id="attributes">Attributes:</h3>
|
||
<pre><code>- file (str): Path/file to config file.
|
||
|
||
- key (str): Path/file to RSA key file.
|
||
|
||
- config (dict): Dictionary containing information of connection
|
||
manager configuration.
|
||
|
||
- connections (dict): Dictionary containing all the nodes added to
|
||
connection manager.
|
||
|
||
- profiles (dict): Dictionary containing all the profiles added to
|
||
connection manager.
|
||
|
||
- privatekey (obj): Object containing the private key to encrypt
|
||
passwords.
|
||
|
||
- publickey (obj): Object containing the public key to decrypt
|
||
passwords.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- conf (str): Path/file to config file. If left empty default
|
||
path is ~/.config/conn/config.yaml
|
||
|
||
- key (str): Path/file to RSA key file. If left empty default
|
||
path is ~/.config/conn/.osk
|
||
</code></pre></div>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="connpy.configfile.encrypt"><code class="name flex">
|
||
<span>def <span class="ident">encrypt</span></span>(<span>self, password, keyfile=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def encrypt(self, password, keyfile=None):
|
||
'''
|
||
Encrypts password using RSA keyfile
|
||
|
||
### Parameters:
|
||
|
||
- password (str): Plaintext password to encrypt.
|
||
|
||
### Optional Parameters:
|
||
|
||
- keyfile (str): Path/file to keyfile. Default is config keyfile.
|
||
|
||
|
||
### Returns:
|
||
|
||
str: Encrypted password.
|
||
|
||
'''
|
||
if keyfile is None:
|
||
keyfile = self.key
|
||
with open(keyfile) as f:
|
||
key = RSA.import_key(f.read())
|
||
f.close()
|
||
publickey = key.publickey()
|
||
encryptor = PKCS1_OAEP.new(publickey)
|
||
password = encryptor.encrypt(password.encode("utf-8"))
|
||
return str(password)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Encrypts password using RSA keyfile</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- password (str): Plaintext password to encrypt.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- keyfile (str): Path/file to keyfile. Default is config keyfile.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>str: Encrypted password.
|
||
</code></pre></div>
|
||
</dd>
|
||
<dt id="connpy.configfile.getitem"><code class="name flex">
|
||
<span>def <span class="ident">getitem</span></span>(<span>self, unique, keys=None, extract=False)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def getitem(self, unique, keys = None, extract = False):
|
||
'''
|
||
Get an node or a group of nodes from configfile which can be passed to node/nodes class
|
||
|
||
### Parameters:
|
||
|
||
- unique (str): Unique name of the node or folder in config using
|
||
connection manager style: node[@subfolder][@folder]
|
||
or [@subfolder]@folder
|
||
|
||
### Optional Parameters:
|
||
|
||
- keys (list): In case you pass a folder as unique, you can filter
|
||
nodes inside the folder passing a list.
|
||
- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
|
||
'''
|
||
uniques = self._explode_unique(unique)
|
||
if unique.startswith("@"):
|
||
if uniques.keys() >= {"folder", "subfolder"}:
|
||
folder = self.connections[uniques["folder"]][uniques["subfolder"]]
|
||
else:
|
||
folder = self.connections[uniques["folder"]]
|
||
newfolder = deepcopy(folder)
|
||
newfolder.pop("type")
|
||
for node_name in folder.keys():
|
||
if node_name == "type":
|
||
continue
|
||
if "type" in newfolder[node_name].keys():
|
||
if newfolder[node_name]["type"] == "subfolder":
|
||
newfolder.pop(node_name)
|
||
else:
|
||
newfolder[node_name].pop("type")
|
||
|
||
if keys != None:
|
||
newfolder = dict((k, newfolder[k]) for k in keys)
|
||
|
||
if extract:
|
||
for node_name, node_keys in newfolder.items():
|
||
for key, value in node_keys.items():
|
||
profile = re.search("^@(.*)", str(value))
|
||
if profile:
|
||
try:
|
||
newfolder[node_name][key] = self.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
newfolder[node_name][key] = ""
|
||
elif value == '' and key == "protocol":
|
||
try:
|
||
newfolder[node_name][key] = self.profiles["default"][key]
|
||
except KeyError:
|
||
newfolder[node_name][key] = "ssh"
|
||
|
||
newfolder = {"{}{}".format(k,unique):v for k,v in newfolder.items()}
|
||
return newfolder
|
||
else:
|
||
if uniques.keys() >= {"folder", "subfolder"}:
|
||
node = self.connections[uniques["folder"]][uniques["subfolder"]][uniques["id"]]
|
||
elif "folder" in uniques.keys():
|
||
node = self.connections[uniques["folder"]][uniques["id"]]
|
||
else:
|
||
node = self.connections[uniques["id"]]
|
||
newnode = deepcopy(node)
|
||
newnode.pop("type")
|
||
|
||
if extract:
|
||
for key, value in newnode.items():
|
||
profile = re.search("^@(.*)", str(value))
|
||
if profile:
|
||
try:
|
||
newnode[key] = self.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
newnode[key] = ""
|
||
elif value == '' and key == "protocol":
|
||
try:
|
||
newnode[key] = self.profiles["default"][key]
|
||
except KeyError:
|
||
newnode[key] = "ssh"
|
||
return newnode</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get an node or a group of nodes from configfile which can be passed to node/nodes class</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- unique (str): Unique name of the node or folder in config using
|
||
connection manager style: node[@subfolder][@folder]
|
||
or [@subfolder]@folder
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- keys (list): In case you pass a folder as unique, you can filter
|
||
nodes inside the folder passing a list.
|
||
- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
</code></pre></div>
|
||
</dd>
|
||
<dt id="connpy.configfile.getitems"><code class="name flex">
|
||
<span>def <span class="ident">getitems</span></span>(<span>self, uniques, extract=False)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def getitems(self, uniques, extract = False):
|
||
'''
|
||
Get a group of nodes from configfile which can be passed to node/nodes class
|
||
|
||
### Parameters:
|
||
|
||
- uniques (str/list): String name that will match hostnames
|
||
from the connection manager. It can be a
|
||
list of strings.
|
||
|
||
### Optional Parameters:
|
||
|
||
- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
|
||
'''
|
||
nodes = {}
|
||
if isinstance(uniques, str):
|
||
uniques = [uniques]
|
||
for i in uniques:
|
||
if i.startswith("@"):
|
||
if not self.config["case"]:
|
||
i = i.lower()
|
||
this = self.getitem(i, extract = extract)
|
||
nodes.update(this)
|
||
else:
|
||
if not self.config["case"]:
|
||
i = i.lower()
|
||
this = self.getitem(i, extract = extract)
|
||
nodes[i] = this
|
||
return nodes</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Get a group of nodes from configfile which can be passed to node/nodes class</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- uniques (str/list): String name that will match hostnames
|
||
from the connection manager. It can be a
|
||
list of strings.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- extract (bool): If True, extract information from profiles.
|
||
Default False.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>dict: Dictionary containing information of node or multiple
|
||
dictionaries of multiple nodes.
|
||
</code></pre></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
<dt id="connpy.node"><code class="flex name class">
|
||
<span>class <span class="ident">node</span></span>
|
||
<span>(</span><span>unique,<br>host,<br>options='',<br>logs='',<br>password='',<br>port='',<br>protocol='',<br>user='',<br>config='',<br>tags='',<br>jumphost='')</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@ClassHook
|
||
class node:
|
||
''' This class generates a node object. Containts all the information and methods to connect and interact with a device using ssh or telnet.
|
||
|
||
### Attributes:
|
||
|
||
- output (str): Output of the commands you ran with run or test
|
||
method.
|
||
|
||
- result(bool): True if expected value is found after running
|
||
the commands using test method.
|
||
|
||
- status (int): 0 if the method run or test run successfully.
|
||
1 if connection failed.
|
||
2 if expect timeouts without prompt or EOF.
|
||
|
||
'''
|
||
|
||
def __init__(self, unique, host, options='', logs='', password='', port='', protocol='', user='', config='', tags='', jumphost=''):
|
||
'''
|
||
|
||
### Parameters:
|
||
|
||
- unique (str): Unique name to assign to the node.
|
||
|
||
- host (str): IP address or hostname of the node.
|
||
|
||
### Optional Parameters:
|
||
|
||
- options (str): Additional options to pass the ssh/telnet for
|
||
connection.
|
||
|
||
- logs (str): Path/file for storing the logs. You can use
|
||
${unique},${host}, ${port}, ${user}, ${protocol}
|
||
as variables.
|
||
|
||
- password (str): Encrypted or plaintext password.
|
||
|
||
- port (str): Port to connect to node, default 22 for ssh and 23
|
||
for telnet.
|
||
|
||
- protocol (str): Select ssh, telnet, kubectl or docker. Default is ssh.
|
||
|
||
- user (str): Username to of the node.
|
||
|
||
- config (obj): Pass the object created with class configfile with
|
||
key for decryption and extra configuration if you
|
||
are using connection manager.
|
||
|
||
- tags (dict) : Tags useful for automation and personal porpuse
|
||
like "os", "prompt" and "screenleght_command"
|
||
|
||
- jumphost (str): Reference another node to be used as a jumphost
|
||
'''
|
||
self.config = config
|
||
if config == '':
|
||
self.idletime = 0
|
||
self.key = None
|
||
else:
|
||
self.idletime = config.config["idletime"]
|
||
self.key = config.key
|
||
self.unique = unique
|
||
attr = {"host": host, "logs": logs, "options":options, "port": port, "protocol": protocol, "user": user, "tags": tags, "jumphost": jumphost}
|
||
for key in attr:
|
||
profile = re.search("^@(.*)", str(attr[key]))
|
||
if profile and config != '':
|
||
try:
|
||
setattr(self,key,config.profiles[profile.group(1)][key])
|
||
except KeyError:
|
||
setattr(self,key,"")
|
||
elif attr[key] == '' and key == "protocol":
|
||
try:
|
||
setattr(self,key,config.profiles["default"][key])
|
||
except (KeyError, AttributeError):
|
||
setattr(self,key,"ssh")
|
||
else:
|
||
setattr(self,key,attr[key])
|
||
if isinstance(password,list):
|
||
self.password = []
|
||
for i, s in enumerate(password):
|
||
profile = re.search("^@(.*)", password[i])
|
||
if profile and config != '':
|
||
self.password.append(config.profiles[profile.group(1)]["password"])
|
||
else:
|
||
self.password.append(password[i])
|
||
else:
|
||
self.password = [password]
|
||
if self.jumphost != "" and config != '':
|
||
self.jumphost = config.getitem(self.jumphost)
|
||
for key in self.jumphost:
|
||
profile = re.search("^@(.*)", str(self.jumphost[key]))
|
||
if profile:
|
||
try:
|
||
self.jumphost[key] = config.profiles[profile.group(1)][key]
|
||
except KeyError:
|
||
self.jumphost[key] = ""
|
||
elif self.jumphost[key] == '' and key == "protocol":
|
||
try:
|
||
self.jumphost[key] = config.profiles["default"][key]
|
||
except KeyError:
|
||
self.jumphost[key] = "ssh"
|
||
if isinstance(self.jumphost["password"],list):
|
||
jumphost_password = []
|
||
for i, s in enumerate(self.jumphost["password"]):
|
||
profile = re.search("^@(.*)", self.jumphost["password"][i])
|
||
if profile:
|
||
jumphost_password.append(config.profiles[profile.group(1)]["password"])
|
||
else:
|
||
jumphost_password.append(self.jumphost["password"][i])
|
||
self.jumphost["password"] = jumphost_password
|
||
else:
|
||
self.jumphost["password"] = [self.jumphost["password"]]
|
||
if self.jumphost["password"] != [""]:
|
||
self.password = self.jumphost["password"] + self.password
|
||
|
||
if self.jumphost["protocol"] == "ssh":
|
||
jumphost_cmd = self.jumphost["protocol"] + " -W %h:%p"
|
||
if self.jumphost["port"] != '':
|
||
jumphost_cmd = jumphost_cmd + " -p " + self.jumphost["port"]
|
||
if self.jumphost["options"] != '':
|
||
jumphost_cmd = jumphost_cmd + " " + self.jumphost["options"]
|
||
if self.jumphost["user"] == '':
|
||
jumphost_cmd = jumphost_cmd + " {}".format(self.jumphost["host"])
|
||
else:
|
||
jumphost_cmd = jumphost_cmd + " {}".format("@".join([self.jumphost["user"],self.jumphost["host"]]))
|
||
self.jumphost = f"-o ProxyCommand=\"{jumphost_cmd}\""
|
||
elif self.jumphost["protocol"] == "ssm":
|
||
ssm_target = self.jumphost["host"]
|
||
ssm_cmd = f"aws ssm start-session --target {ssm_target} --document-name AWS-StartSSHSession --parameters 'portNumber=22'"
|
||
if isinstance(self.jumphost.get("tags"), dict):
|
||
if "profile" in self.jumphost["tags"]:
|
||
ssm_cmd += f" --profile {self.jumphost['tags']['profile']}"
|
||
if "region" in self.jumphost["tags"]:
|
||
ssm_cmd += f" --region {self.jumphost['tags']['region']}"
|
||
if self.jumphost["options"] != '':
|
||
ssm_cmd += f" {self.jumphost['options']}"
|
||
|
||
bastion_user_part = f"{self.jumphost['user']}@{ssm_target}" if self.jumphost['user'] else ssm_target
|
||
|
||
ssh_opts = ""
|
||
if isinstance(self.jumphost.get("tags"), dict) and "ssh_options" in self.jumphost["tags"]:
|
||
ssh_opts = f" {self.jumphost['tags']['ssh_options']}"
|
||
|
||
inner_ssh = f"ssh{ssh_opts} -o ProxyCommand='{ssm_cmd}' -W %h:%p {bastion_user_part}"
|
||
self.jumphost = f"-o ProxyCommand=\"{inner_ssh}\""
|
||
elif self.jumphost["protocol"] in ["kubectl", "docker"]:
|
||
nc_cmd = "nc"
|
||
if isinstance(self.jumphost.get("tags"), dict) and "nc_command" in self.jumphost["tags"]:
|
||
nc_cmd = self.jumphost["tags"]["nc_command"]
|
||
|
||
if self.jumphost["protocol"] == "kubectl":
|
||
proxy_cmd = f"kubectl exec "
|
||
if self.jumphost["options"] != '':
|
||
proxy_cmd += f"{self.jumphost['options']} "
|
||
proxy_cmd += f"{self.jumphost['host']} -i -- {nc_cmd} %h %p"
|
||
else:
|
||
proxy_cmd = f"docker "
|
||
if self.jumphost["options"] != '':
|
||
proxy_cmd += f"{self.jumphost['options']} "
|
||
proxy_cmd += f"exec -i {self.jumphost['host']} {nc_cmd} %h %p"
|
||
|
||
self.jumphost = f"-o ProxyCommand=\"{proxy_cmd}\""
|
||
else:
|
||
self.jumphost = ""
|
||
|
||
self.output = ""
|
||
self.status = 1
|
||
self.result = {}
|
||
|
||
@MethodHook
|
||
def _passtx(self, passwords, *, keyfile=None):
|
||
# decrypts passwords, used by other methdos.
|
||
dpass = []
|
||
if keyfile is None:
|
||
keyfile = self.key
|
||
if keyfile is not None:
|
||
with open(keyfile) as f:
|
||
key = RSA.import_key(f.read())
|
||
decryptor = PKCS1_OAEP.new(key)
|
||
for passwd in passwords:
|
||
if not re.match('^b[\"\'].+[\"\']$', passwd):
|
||
dpass.append(passwd)
|
||
else:
|
||
try:
|
||
decrypted = decryptor.decrypt(ast.literal_eval(passwd)).decode("utf-8")
|
||
dpass.append(decrypted)
|
||
except Exception:
|
||
printer.error("Decryption failed: Missing or corrupted key.")
|
||
printer.info("Verify your RSA key and configuration settings.")
|
||
sys.exit(1)
|
||
return dpass
|
||
|
||
|
||
|
||
@MethodHook
|
||
def _logfile(self, logfile = None):
|
||
# translate logs variables and generate logs path.
|
||
if logfile == None:
|
||
logfile = self.logs
|
||
logfile = logfile.replace("${unique}", self.unique)
|
||
logfile = logfile.replace("${host}", self.host)
|
||
logfile = logfile.replace("${port}", self.port)
|
||
logfile = logfile.replace("${user}", self.user)
|
||
logfile = logfile.replace("${protocol}", self.protocol)
|
||
now = datetime.datetime.now()
|
||
dateconf = re.search(r'\$\{date \'(.*)\'}', logfile)
|
||
if dateconf:
|
||
logfile = re.sub(r'\$\{date (.*)}',now.strftime(dateconf.group(1)), logfile)
|
||
return logfile
|
||
|
||
@MethodHook
|
||
def _logclean(self, logfile, var = False):
|
||
"""Remove special ascii characters and process terminal cursor movements to clean logs."""
|
||
from .utils import log_cleaner
|
||
|
||
if var == False:
|
||
try:
|
||
with open(logfile, "r") as f:
|
||
t = f.read()
|
||
except:
|
||
return
|
||
else:
|
||
t = logfile
|
||
|
||
result = log_cleaner(t)
|
||
|
||
if var == False:
|
||
try:
|
||
with open(logfile, "w") as f:
|
||
f.write(result)
|
||
except:
|
||
pass
|
||
return
|
||
else:
|
||
return result
|
||
|
||
@MethodHook
|
||
def _savelog(self):
|
||
'''Save the log buffer to the file at regular intervals if there are changes.'''
|
||
t = threading.current_thread()
|
||
prev_size = 0 # Store the previous size of the buffer
|
||
|
||
while getattr(t, "do_run", True): # Check if thread is signaled to stop
|
||
current_size = self.mylog.tell() # Current size of the buffer
|
||
|
||
# Only save if the buffer size has changed
|
||
if current_size != prev_size:
|
||
with open(self.logfile, "w") as f: # Use "w" to overwrite the file
|
||
f.write(self._logclean(self.mylog.getvalue().decode(), True))
|
||
prev_size = current_size # Update the previous size
|
||
sleep(5)
|
||
|
||
@MethodHook
|
||
def _filter(self, a):
|
||
#Set time for last input when using interact
|
||
self.lastinput = time()
|
||
return a
|
||
|
||
@MethodHook
|
||
def _keepalive(self):
|
||
#Send keepalive ctrl+e when idletime passed without new inputs on interact
|
||
self.lastinput = time()
|
||
t = threading.current_thread()
|
||
while True:
|
||
if time() - self.lastinput >= self.idletime:
|
||
self.child.sendcontrol("e")
|
||
self.lastinput = time()
|
||
sleep(1)
|
||
|
||
|
||
def _setup_interact_environment(self, debug=False, logger=None, async_mode=False):
|
||
size = re.search('columns=([0-9]+).*lines=([0-9]+)',str(os.get_terminal_size()))
|
||
self.child.setwinsize(int(size.group(2)),int(size.group(1)))
|
||
if logger:
|
||
port_str = f":{self.port}" if self.port and self.protocol not in ["ssm", "kubectl", "docker"] else ""
|
||
logger("success", f"Connected to {self.unique} at {self.host}{port_str} via: {self.protocol}")
|
||
|
||
# Always initialize self.mylog to capture terminal context for the AI Copilot
|
||
if not hasattr(self, 'mylog'):
|
||
self.mylog = io.BytesIO()
|
||
|
||
if not async_mode:
|
||
self.child.logfile_read = self.mylog
|
||
|
||
# Only start disk-logging tasks if logfile is configured
|
||
if 'logfile' in dir(self):
|
||
if not async_mode:
|
||
# Start the _savelog thread (sync mode)
|
||
log_thread = threading.Thread(target=self._savelog)
|
||
log_thread.daemon = True
|
||
log_thread.start()
|
||
if 'missingtext' in dir(self):
|
||
print(self.child.after.decode(), end='')
|
||
if self.idletime > 0 and not async_mode:
|
||
x = threading.Thread(target=self._keepalive)
|
||
x.daemon = True
|
||
x.start()
|
||
if debug:
|
||
if 'mylog' in dir(self):
|
||
if not async_mode:
|
||
print(self.mylog.getvalue().decode())
|
||
|
||
def _teardown_interact_environment(self):
|
||
if 'logfile' in dir(self) and hasattr(self, 'mylog'):
|
||
with open(self.logfile, "w") as f:
|
||
f.write(self._logclean(self.mylog.getvalue().decode(), True))
|
||
|
||
async def _async_interact_loop(self, local_stream, resize_callback, copilot_handler=None):
|
||
local_stream.setup(resize_callback=resize_callback)
|
||
try:
|
||
child_fd = self.child.child_fd
|
||
|
||
# 1. Flush ghost buffer (Clean UX)
|
||
ghost_buffer = b''
|
||
if getattr(self, 'missingtext', False):
|
||
# If we are missing the password, we MUST show the password prompt
|
||
ghost_buffer = (self.child.after or b'') + (self.child.buffer or b'')
|
||
else:
|
||
# We auto-logged in. Hide the messy password negotiation and just keep any pending live stream.
|
||
ghost_buffer = self.child.buffer or b''
|
||
|
||
# Fix user's pet peeve: Strip leading newlines to avoid the empty lines
|
||
# the router echoes after receiving the password or blank line.
|
||
if not getattr(self, 'missingtext', False):
|
||
ghost_buffer = ghost_buffer.lstrip(b'\r\n ')
|
||
|
||
if ghost_buffer:
|
||
# Add a single clean newline so it doesn't merge with the Connected message
|
||
await local_stream.write(b'\r\n' + ghost_buffer)
|
||
if hasattr(self, 'mylog'):
|
||
self.mylog.write(b'\n' + ghost_buffer)
|
||
|
||
self.child.buffer = b''
|
||
self.child.before = b''
|
||
|
||
# 2. Set child fd non-blocking
|
||
flags = fcntl.fcntl(child_fd, fcntl.F_GETFL)
|
||
fcntl.fcntl(child_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK)
|
||
|
||
loop = asyncio.get_running_loop()
|
||
child_reader_queue = asyncio.Queue()
|
||
|
||
# Track command byte positions for copilot context navigation
|
||
# Each entry is (byte_position, command_text_or_None)
|
||
cmd_byte_positions = [(0, None)]
|
||
|
||
def _child_read_ready():
|
||
try:
|
||
# Increase buffer to 64KB for better high-speed handling
|
||
data = os.read(child_fd, 65536)
|
||
if data:
|
||
child_reader_queue.put_nowait(data)
|
||
else:
|
||
child_reader_queue.put_nowait(b'')
|
||
except BlockingIOError:
|
||
pass
|
||
except OSError:
|
||
child_reader_queue.put_nowait(b'')
|
||
|
||
loop.add_reader(child_fd, _child_read_ready)
|
||
self.lastinput = time()
|
||
|
||
async def ingress_task():
|
||
while True:
|
||
data = await local_stream.read()
|
||
if not data:
|
||
break
|
||
|
||
# Copilot interception
|
||
if copilot_handler and b'\x00' in data:
|
||
# Build node info from available metadata and ensure values are strings (not bytes)
|
||
def to_str(val):
|
||
if isinstance(val, bytes):
|
||
return val.decode(errors='replace')
|
||
return str(val) if val is not None else "unknown"
|
||
|
||
node_info = {
|
||
"name": to_str(getattr(self, 'unique', 'unknown')),
|
||
"host": to_str(getattr(self, 'host', 'unknown'))
|
||
}
|
||
if isinstance(getattr(self, 'tags', None), dict):
|
||
node_info["os"] = to_str(self.tags.get("os", "unknown"))
|
||
node_info["prompt"] = to_str(self.tags.get("prompt", r'>$|#$|\$$|>.$|#.$|\$.$'))
|
||
|
||
# Invoke copilot (async callback handles UI)
|
||
await copilot_handler(self.mylog.getvalue(), node_info, local_stream, child_fd, cmd_byte_positions)
|
||
continue
|
||
|
||
# Remove any stray \x00 bytes and forward normally
|
||
clean_data = data.replace(b'\x00', b'')
|
||
if clean_data:
|
||
# Track command boundaries when user hits Enter
|
||
if hasattr(self, 'mylog') and (b'\r' in clean_data or b'\n' in clean_data):
|
||
cmd_byte_positions.append((self.mylog.tell(), None))
|
||
|
||
try:
|
||
os.write(child_fd, clean_data)
|
||
except OSError:
|
||
break
|
||
self.lastinput = time()
|
||
|
||
async def egress_task():
|
||
# Continue stripping newlines from the live stream until we hit real text
|
||
skip_newlines = not getattr(self, 'missingtext', False) and not ghost_buffer
|
||
while True:
|
||
data = await child_reader_queue.get()
|
||
if not data:
|
||
break
|
||
|
||
# Batching Optimization: Drain the queue to batch writes during high-volume bursts
|
||
# Helps the terminal parse ANSI faster and reduces syscalls.
|
||
chunks = [data]
|
||
while not child_reader_queue.empty():
|
||
try:
|
||
extra = child_reader_queue.get_nowait()
|
||
if not extra:
|
||
chunks.append(b'') # Re-put EOF later or handle it
|
||
break
|
||
chunks.append(extra)
|
||
except asyncio.QueueEmpty:
|
||
break
|
||
|
||
has_eof = chunks[-1] == b''
|
||
if has_eof:
|
||
chunks.pop()
|
||
|
||
if chunks:
|
||
combined_data = b''.join(chunks)
|
||
if skip_newlines:
|
||
stripped = combined_data.lstrip(b'\r\n')
|
||
if stripped:
|
||
skip_newlines = False
|
||
combined_data = stripped
|
||
else:
|
||
if has_eof: break
|
||
continue
|
||
|
||
await local_stream.write(combined_data)
|
||
if hasattr(self, 'mylog'):
|
||
self.mylog.write(combined_data)
|
||
|
||
if has_eof:
|
||
break
|
||
|
||
async def keepalive_task():
|
||
while True:
|
||
await asyncio.sleep(1)
|
||
if time() - self.lastinput >= self.idletime:
|
||
try:
|
||
self.child.sendcontrol("e")
|
||
self.lastinput = time()
|
||
except Exception:
|
||
pass
|
||
|
||
async def savelog_task():
|
||
prev_size = 0
|
||
while True:
|
||
await asyncio.sleep(5)
|
||
current_size = self.mylog.tell()
|
||
if current_size != prev_size:
|
||
try:
|
||
# Move heavy log cleaning to a thread to avoid freezing the interaction loop
|
||
raw_log = self.mylog.getvalue().decode(errors='replace')
|
||
cleaned_log = await asyncio.to_thread(self._logclean, raw_log, True)
|
||
with open(self.logfile, "w") as f:
|
||
f.write(cleaned_log)
|
||
prev_size = current_size
|
||
except Exception:
|
||
pass
|
||
|
||
try:
|
||
# We wait for either the user (ingress) or the child (egress) to finish
|
||
tasks = [
|
||
asyncio.create_task(ingress_task()),
|
||
asyncio.create_task(egress_task())
|
||
]
|
||
if self.idletime > 0:
|
||
tasks.append(asyncio.create_task(keepalive_task()))
|
||
if hasattr(self, 'logfile') and hasattr(self, 'mylog'):
|
||
tasks.append(asyncio.create_task(savelog_task()))
|
||
|
||
done, pending = await asyncio.wait(
|
||
[tasks[0], tasks[1]],
|
||
return_when=asyncio.FIRST_COMPLETED
|
||
)
|
||
|
||
# If ingress finished first (user quit), give egress a small window to catch up
|
||
# on the remaining output in the queue.
|
||
if tasks[0] in done and tasks[1] not in done:
|
||
try:
|
||
await asyncio.wait_for(tasks[1], timeout=0.2)
|
||
except (asyncio.TimeoutError, asyncio.CancelledError):
|
||
pass
|
||
|
||
for t in tasks:
|
||
if t not in done:
|
||
t.cancel()
|
||
|
||
# Final log sync on thread to avoid losing last lines
|
||
if hasattr(self, 'logfile') and hasattr(self, 'mylog'):
|
||
try:
|
||
raw_log = self.mylog.getvalue().decode(errors='replace')
|
||
cleaned_log = await asyncio.to_thread(self._logclean, raw_log, True)
|
||
with open(self.logfile, "w") as f:
|
||
f.write(cleaned_log)
|
||
except Exception:
|
||
pass
|
||
|
||
finally:
|
||
loop.remove_reader(child_fd)
|
||
try:
|
||
flags = fcntl.fcntl(child_fd, fcntl.F_GETFL)
|
||
fcntl.fcntl(child_fd, fcntl.F_SETFL, flags & ~os.O_NONBLOCK)
|
||
except Exception:
|
||
pass
|
||
finally:
|
||
local_stream.teardown()
|
||
|
||
|
||
@MethodHook
|
||
def interact(self, debug=False, logger=None):
|
||
'''
|
||
Asynchronous interactive session using Smart Tunnel architecture.
|
||
Allows multiplexing I/O and handling SIGWINCH events locally without blocking.
|
||
'''
|
||
connect = self._connect(debug=debug, logger=logger)
|
||
if connect == True:
|
||
try:
|
||
self._setup_interact_environment(debug=debug, logger=logger, async_mode=True)
|
||
|
||
local_stream = LocalStream()
|
||
|
||
def resize_callback(rows, cols):
|
||
try:
|
||
self.child.setwinsize(rows, cols)
|
||
except Exception:
|
||
pass
|
||
|
||
# Build local copilot handler
|
||
copilot_handler = self._build_local_copilot_handler()
|
||
|
||
asyncio.run(self._async_interact_loop(local_stream, resize_callback, copilot_handler=copilot_handler))
|
||
finally:
|
||
self._teardown_interact_environment()
|
||
else:
|
||
if logger:
|
||
logger("error", str(connect))
|
||
else:
|
||
printer.error(f"Connection failed: {str(connect)}")
|
||
sys.exit(1)
|
||
|
||
def _build_local_copilot_handler(self):
|
||
"""Build copilot handler for local CLI sessions using rich for rendering."""
|
||
config = getattr(self, 'config', None) if hasattr(self, 'config') else None
|
||
return self._copilot_handler(config)
|
||
|
||
def _copilot_handler(self, config):
|
||
"""Unified copilot handler for local session."""
|
||
from .cli.terminal_ui import CopilotInterface
|
||
from .services.ai_service import AIService
|
||
import asyncio
|
||
import os
|
||
|
||
async def handler(buffer, node_info, stream, child_fd, cmd_byte_positions=None):
|
||
try:
|
||
interface = CopilotInterface(
|
||
config,
|
||
history=getattr(stream, 'copilot_history', None),
|
||
session_state=getattr(stream, 'copilot_state', None)
|
||
)
|
||
# Save history back to stream for persistence in current session
|
||
stream.copilot_history = interface.history
|
||
stream.copilot_state = interface.session_state
|
||
|
||
ai_service = AIService(config)
|
||
|
||
async def on_ai_call(active_buffer, question, chunk_callback, merged_node_info):
|
||
return await ai_service.aask_copilot(
|
||
active_buffer,
|
||
question,
|
||
node_info=merged_node_info,
|
||
chunk_callback=chunk_callback
|
||
)
|
||
# Get raw bytes from BytesIO
|
||
raw_bytes = self.mylog.getvalue()
|
||
|
||
# Detener el lector de la terminal para que prompt_toolkit (en run_session)
|
||
# tenga control exclusivo del stdin sin interferencias de LocalStream.
|
||
if hasattr(stream, 'stop_reading'):
|
||
stream.stop_reading()
|
||
elif hasattr(stream, '_loop') and hasattr(stream, 'stdin_fd'):
|
||
# Fallback si no tiene el método (en LocalStream)
|
||
stream._loop.remove_reader(stream.stdin_fd)
|
||
|
||
try:
|
||
with copilot_terminal_mode():
|
||
while True:
|
||
action, commands, custom_cmd = await interface.run_session(
|
||
raw_bytes=raw_bytes,
|
||
cmd_byte_positions=cmd_byte_positions,
|
||
node_info=node_info,
|
||
on_ai_call=on_ai_call
|
||
)
|
||
if action == "continue":
|
||
continue
|
||
break
|
||
finally:
|
||
# Reiniciar el lector de la terminal para volver al modo interactivo SSH/Telnet
|
||
if hasattr(stream, 'start_reading'):
|
||
stream.start_reading()
|
||
elif hasattr(stream, '_loop') and hasattr(stream, 'stdin_fd'):
|
||
stream._loop.add_reader(stream.stdin_fd, stream._read_ready)
|
||
|
||
if action in ("send_all", "custom"):
|
||
cmds_to_send = commands if action == "send_all" else custom_cmd
|
||
|
||
if cmds_to_send:
|
||
os.write(child_fd, b'\x15') # Ctrl+U
|
||
await asyncio.sleep(0.1)
|
||
|
||
# Prepend screen length command to avoid pagination
|
||
if "screen_length_command" in self.tags:
|
||
cmds_to_send.insert(0, self.tags["screen_length_command"])
|
||
|
||
for cmd in cmds_to_send:
|
||
if cmd_byte_positions is not None:
|
||
cmd_byte_positions.append((self.mylog.tell(), cmd))
|
||
os.write(child_fd, (cmd + "\n").encode())
|
||
await asyncio.sleep(0.8)
|
||
else:
|
||
os.write(child_fd, b'\x15\r')
|
||
except Exception as e:
|
||
import traceback
|
||
print(f"\n[ERROR in Copilot Handler] {e}", flush=True)
|
||
traceback.print_exc()
|
||
os.write(child_fd, b'\x15\r')
|
||
|
||
return handler
|
||
|
||
@MethodHook
|
||
def run(self, commands, vars = None,*, folder = '', prompt = r'>$|#$|\$$|>.$|#.$|\$.$', stdout = False, timeout = 10, logger = None):
|
||
'''
|
||
Run a command or list of commands on the node and return the output.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables
|
||
used in commands parameter.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave
|
||
empty to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- stdout (bool):Set True to send the command output to stdout.
|
||
default False.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
### Returns:
|
||
|
||
str: Output of the commands you ran on the node.
|
||
|
||
'''
|
||
connect = self._connect(timeout = timeout, logger = logger)
|
||
now = datetime.datetime.now().strftime('%Y-%m-%d_%H%M%S')
|
||
if connect == True:
|
||
if logger:
|
||
port_str = f":{self.port}" if self.port and self.protocol not in ["ssm", "kubectl", "docker"] else ""
|
||
logger("success", f"Connected to {self.unique} at {self.host}{port_str} via: {self.protocol}")
|
||
|
||
# Attempt to set the terminal size
|
||
try:
|
||
self.child.setwinsize(65535, 65535)
|
||
except Exception:
|
||
try:
|
||
self.child.setwinsize(10000, 10000)
|
||
except Exception:
|
||
pass
|
||
if "prompt" in self.tags:
|
||
prompt = self.tags["prompt"]
|
||
expects = [prompt, pexpect.EOF, pexpect.TIMEOUT]
|
||
output = ''
|
||
status = ''
|
||
if not isinstance(commands, list):
|
||
commands = [commands]
|
||
if "screen_length_command" in self.tags:
|
||
commands.insert(0, self.tags["screen_length_command"])
|
||
self.mylog = io.BytesIO()
|
||
self.child.logfile_read = self.mylog
|
||
for c in commands:
|
||
if vars is not None:
|
||
try:
|
||
c = c.format(**vars)
|
||
except KeyError as e:
|
||
self.output = f"Error: Variable {e} not defined in task or inventory"
|
||
self.status = 1
|
||
return self.output
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.sendline(c)
|
||
if result == 2:
|
||
break
|
||
if not result == 2:
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.close()
|
||
output = self._logclean(self.mylog.getvalue().decode(), True)
|
||
if logger:
|
||
logger("output", output)
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(output)
|
||
f.close()
|
||
self.output = output
|
||
if result == 2:
|
||
self.status = 2
|
||
else:
|
||
self.status = 0
|
||
return output
|
||
else:
|
||
self.output = connect
|
||
self.status = 1
|
||
if logger:
|
||
logger("error", f"Connection failed: {connect}")
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(connect)
|
||
|
||
f.close()
|
||
return connect
|
||
|
||
@MethodHook
|
||
def test(self, commands, expected, vars = None,*, folder = '', prompt = r'>$|#$|\$$|>.$|#.$|\$.$', timeout = 10, logger = None):
|
||
'''
|
||
Run a command or list of commands on the node, then check if expected value appears on the output after the last command.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
|
||
- expected (str) : Expected text to appear after running
|
||
all the commands on the node.You can use
|
||
variables as {varname} and defining them
|
||
in optional parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables
|
||
used in commands and expected parameters.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave
|
||
empty to not store logs.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
### Returns:
|
||
bool: true if expected value is found after running the commands
|
||
false if prompt is found before.
|
||
|
||
'''
|
||
now = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
connect = self._connect(timeout = timeout, logger = logger)
|
||
if connect == True:
|
||
if logger:
|
||
port_str = f":{self.port}" if self.port and self.protocol not in ["ssm", "kubectl", "docker"] else ""
|
||
logger("success", f"Connected to {self.unique} at {self.host}{port_str} via: {self.protocol}")
|
||
|
||
# Attempt to set the terminal size
|
||
try:
|
||
self.child.setwinsize(65535, 65535)
|
||
except Exception:
|
||
try:
|
||
self.child.setwinsize(10000, 10000)
|
||
except Exception:
|
||
pass
|
||
if "prompt" in self.tags:
|
||
prompt = self.tags["prompt"]
|
||
expects = [prompt, pexpect.EOF, pexpect.TIMEOUT]
|
||
output = ''
|
||
if not isinstance(commands, list):
|
||
commands = [commands]
|
||
if not isinstance(expected, list):
|
||
expected = [expected]
|
||
if "screen_length_command" in self.tags:
|
||
commands.insert(0, self.tags["screen_length_command"])
|
||
self.mylog = io.BytesIO()
|
||
self.child.logfile_read = self.mylog
|
||
for c in commands:
|
||
if vars is not None:
|
||
try:
|
||
c = c.format(**vars)
|
||
except KeyError as e:
|
||
self.output = f"Error: Variable {e} not defined in task or inventory"
|
||
self.status = 1
|
||
return self.output
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.sendline(c)
|
||
if result == 2:
|
||
break
|
||
if not result == 2:
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.close()
|
||
output = self._logclean(self.mylog.getvalue().decode(), True)
|
||
if logger:
|
||
logger("output", output)
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(output)
|
||
f.close()
|
||
self.output = output
|
||
if result in [0, 1]:
|
||
# lastcommand = commands[-1]
|
||
# if vars is not None:
|
||
# lastcommand = lastcommand.format(**vars)
|
||
# last_command_index = output.rfind(lastcommand)
|
||
# cleaned_output = output[last_command_index + len(lastcommand):].strip()
|
||
self.result = {}
|
||
for e in expected:
|
||
if vars is not None:
|
||
e = e.format(**vars)
|
||
updatedprompt = re.sub(r'(?<!\\)\$', '', prompt)
|
||
newpattern = f".*({updatedprompt}).*{e}.*"
|
||
cleaned_output = output
|
||
cleaned_output = re.sub(newpattern, '', cleaned_output)
|
||
if e in cleaned_output:
|
||
self.result[e] = True
|
||
else:
|
||
self.result[e]= False
|
||
self.status = 0
|
||
return self.result
|
||
if result == 2:
|
||
self.result = None
|
||
self.status = 2
|
||
return output
|
||
else:
|
||
self.result = None
|
||
self.output = connect
|
||
self.status = 1
|
||
return connect
|
||
|
||
@MethodHook
|
||
def _generate_ssh_sftp_cmd(self):
|
||
cmd = self.protocol
|
||
if self.port:
|
||
if self.protocol == "ssh":
|
||
cmd += " -p " + self.port
|
||
elif self.protocol == "sftp":
|
||
cmd += " -P " + self.port
|
||
if self.options:
|
||
opts = self.options
|
||
if self.protocol == "sftp":
|
||
# Strip SSH-only flags that sftp doesn't support
|
||
opts = re.sub(r'(?<!\S)-[XxtTAaNf]\b', '', opts).strip()
|
||
if opts:
|
||
cmd += " " + opts
|
||
if self.jumphost:
|
||
cmd += " " + self.jumphost
|
||
user_host = f"{self.user}@{self.host}" if self.user else self.host
|
||
cmd += f" {user_host}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _generate_telnet_cmd(self):
|
||
cmd = f"telnet {self.host}"
|
||
if self.port:
|
||
cmd += f" {self.port}"
|
||
if self.options:
|
||
cmd += f" {self.options}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _generate_kube_cmd(self):
|
||
cmd = f"kubectl exec {self.options} {self.host} -it --"
|
||
kube_command = self.tags.get("kube_command", "/bin/bash") if isinstance(self.tags, dict) else "/bin/bash"
|
||
cmd += f" {kube_command}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _generate_docker_cmd(self):
|
||
cmd = f"docker {self.options} exec -it {self.host}"
|
||
docker_command = self.tags.get("docker_command", "/bin/bash") if isinstance(self.tags, dict) else "/bin/bash"
|
||
cmd += f" {docker_command}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _generate_ssm_cmd(self):
|
||
region = self.tags.get("region", "") if isinstance(self.tags, dict) else ""
|
||
profile = self.tags.get("profile", "") if isinstance(self.tags, dict) else ""
|
||
cmd = f"aws ssm start-session --target {self.host}"
|
||
if region:
|
||
cmd += f" --region {region}"
|
||
if profile:
|
||
cmd += f" --profile {profile}"
|
||
if self.options:
|
||
cmd += f" {self.options}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _generate_ssm_cmd(self):
|
||
region = self.tags.get("region", "") if isinstance(self.tags, dict) else ""
|
||
profile = self.tags.get("profile", "") if isinstance(self.tags, dict) else ""
|
||
cmd = f"aws ssm start-session --target {self.host}"
|
||
if region:
|
||
cmd += f" --region {region}"
|
||
if profile:
|
||
cmd += f" --profile {profile}"
|
||
if self.options:
|
||
cmd += f" {self.options}"
|
||
return cmd
|
||
|
||
@MethodHook
|
||
def _get_cmd(self):
|
||
if self.protocol in ["ssh", "sftp"]:
|
||
return self._generate_ssh_sftp_cmd()
|
||
elif self.protocol == "telnet":
|
||
return self._generate_telnet_cmd()
|
||
elif self.protocol == "kubectl":
|
||
return self._generate_kube_cmd()
|
||
elif self.protocol == "docker":
|
||
return self._generate_docker_cmd()
|
||
elif self.protocol == "ssm":
|
||
return self._generate_ssm_cmd()
|
||
else:
|
||
printer.error(f"Invalid protocol: {self.protocol}")
|
||
sys.exit(1)
|
||
|
||
@MethodHook
|
||
def _connect(self, debug=False, timeout=10, max_attempts=3, logger=None):
|
||
|
||
cmd = self._get_cmd()
|
||
passwords = self._passtx(self.password) if self.password and any(self.password) else []
|
||
if self.logs != '':
|
||
self.logfile = self._logfile()
|
||
default_prompt = r'>$|#$|\$$|>.$|#.$|\$.$'
|
||
prompt = self.tags.get("prompt", default_prompt) if isinstance(self.tags, dict) else default_prompt
|
||
password_prompt = '[p|P]assword:|[u|U]sername:' if self.protocol != 'telnet' else '[p|P]assword:'
|
||
|
||
expects = {
|
||
"ssh": ['yes/no', 'refused', 'supported', 'Invalid|[u|U]sage: ssh', 'ssh-keygen.*\"', 'timeout|timed.out', 'unavailable', 'closed', password_prompt, prompt, 'suspend', pexpect.EOF, pexpect.TIMEOUT, "No route to host", "resolve hostname", "no matching", "[b|B]ad (owner|permissions)"],
|
||
"sftp": ['yes/no', 'refused', 'supported', 'Invalid|[u|U]sage: sftp', 'ssh-keygen.*\"', 'timeout|timed.out', 'unavailable', 'closed', password_prompt, prompt, 'suspend', pexpect.EOF, pexpect.TIMEOUT, "No route to host", "resolve hostname", "no matching", "[b|B]ad (owner|permissions)"],
|
||
"telnet": ['[u|U]sername:', 'refused', 'supported', 'invalid|unrecognized option', 'ssh-keygen.*\"', 'timeout|timed.out', 'unavailable', 'closed', password_prompt, prompt, 'suspend', pexpect.EOF, pexpect.TIMEOUT, "No route to host", "resolve hostname", "no matching", "[b|B]ad (owner|permissions)"],
|
||
"kubectl": ['[u|U]sername:', '[r|R]efused', '[E|e]rror', 'DEPRECATED', pexpect.TIMEOUT, password_prompt, prompt, pexpect.EOF, "expired|invalid"],
|
||
"docker": ['[u|U]sername:', 'Cannot', '[E|e]rror', 'failed', 'not a docker command', 'unknown', 'unable to resolve', pexpect.TIMEOUT, password_prompt, prompt, pexpect.EOF],
|
||
"ssm": ['[u|U]sername:', 'Cannot', '[E|e]rror', 'failed', 'SessionManagerPlugin', '[u|U]nknown', 'unable to resolve', pexpect.TIMEOUT, password_prompt, prompt, pexpect.EOF]
|
||
}
|
||
|
||
error_indices = {
|
||
"ssh": [1, 2, 3, 4, 5, 6, 7, 12, 13, 14, 15, 16],
|
||
"sftp": [1, 2, 3, 4, 5, 6, 7, 12, 13, 14, 15, 16],
|
||
"telnet": [1, 2, 3, 4, 5, 6, 7, 12, 13, 14, 15, 16],
|
||
"kubectl": [1, 2, 3, 4, 8], # Define error indices for kube
|
||
"docker": [1, 2, 3, 4, 5, 6, 7], # Define error indices for docker
|
||
"ssm": [1, 2, 3, 4, 5, 6, 7]
|
||
}
|
||
|
||
eof_indices = {
|
||
"ssh": [8, 9, 10, 11],
|
||
"sftp": [8, 9, 10, 11],
|
||
"telnet": [8, 9, 10, 11],
|
||
"kubectl": [5, 6, 7], # Define eof indices for kube
|
||
"docker": [8, 9, 10], # Define eof indices for docker
|
||
"ssm": [8, 9, 10]
|
||
}
|
||
|
||
initial_indices = {
|
||
"ssh": [0],
|
||
"sftp": [0],
|
||
"telnet": [0],
|
||
"kubectl": [0], # Define special indices for kube
|
||
"docker": [0], # Define special indices for docker
|
||
"ssm": [0]
|
||
}
|
||
|
||
attempts = 1
|
||
while attempts <= max_attempts:
|
||
child = pexpect.spawn(cmd)
|
||
if isinstance(self.tags, dict) and self.tags.get("console"):
|
||
child.sendline()
|
||
if debug:
|
||
if logger:
|
||
logger("debug", f"Command:\n{cmd}")
|
||
self.mylog = io.BytesIO()
|
||
self.mylog.write(f"[i] [DEBUG] Command:\r\n {cmd}\r\n".encode())
|
||
child.logfile_read = self.mylog
|
||
|
||
|
||
endloop = False
|
||
for i in range(len(passwords) if passwords else 1):
|
||
while True:
|
||
results = child.expect(expects[self.protocol], timeout=timeout)
|
||
results_value = expects[self.protocol][results]
|
||
|
||
if results in initial_indices[self.protocol]:
|
||
if self.protocol in ["ssh", "sftp"]:
|
||
child.sendline('yes')
|
||
elif self.protocol in ["telnet", "kubectl", "docker", "ssm"]:
|
||
if self.user:
|
||
child.sendline(self.user)
|
||
else:
|
||
self.missingtext = True
|
||
break
|
||
|
||
elif results in error_indices[self.protocol]:
|
||
child.terminate()
|
||
if results_value == pexpect.TIMEOUT and attempts != max_attempts:
|
||
attempts += 1
|
||
endloop = True
|
||
break
|
||
else:
|
||
after = "Connection timeout" if results_value == pexpect.TIMEOUT else child.after.decode()
|
||
return f"Connection failed code: {results}\n{child.before.decode().lstrip()}{after}{child.readline().decode()}".rstrip()
|
||
|
||
elif results in eof_indices[self.protocol]:
|
||
if results_value == password_prompt:
|
||
if passwords:
|
||
child.sendline(passwords[i])
|
||
else:
|
||
self.missingtext = True
|
||
break
|
||
elif results_value == "suspend":
|
||
child.sendline("\r")
|
||
sleep(2)
|
||
else:
|
||
endloop = True
|
||
child.sendline()
|
||
break
|
||
|
||
if endloop:
|
||
break
|
||
if results_value == pexpect.TIMEOUT:
|
||
continue
|
||
else:
|
||
break
|
||
|
||
if isinstance(self.tags, dict) and self.tags.get("post_connect_commands"):
|
||
cmds = self.tags.get("post_connect_commands")
|
||
commands = [cmds] if isinstance(cmds, str) else cmds
|
||
for command in commands:
|
||
child.sendline(command)
|
||
sleep(1)
|
||
child.readline(0)
|
||
self.child = child
|
||
from pexpect import fdpexpect
|
||
self.raw_child = fdpexpect.fdspawn(self.child.child_fd)
|
||
return True</code></pre>
|
||
</details>
|
||
<div class="desc"><p>This class generates a node object. Containts all the information and methods to connect and interact with a device using ssh or telnet.</p>
|
||
<h3 id="attributes">Attributes:</h3>
|
||
<pre><code>- output (str): Output of the commands you ran with run or test
|
||
method.
|
||
|
||
- result(bool): True if expected value is found after running
|
||
the commands using test method.
|
||
|
||
- status (int): 0 if the method run or test run successfully.
|
||
1 if connection failed.
|
||
2 if expect timeouts without prompt or EOF.
|
||
</code></pre>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- unique (str): Unique name to assign to the node.
|
||
|
||
- host (str): IP address or hostname of the node.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- options (str): Additional options to pass the ssh/telnet for
|
||
connection.
|
||
|
||
- logs (str): Path/file for storing the logs. You can use
|
||
${unique},${host}, ${port}, ${user}, ${protocol}
|
||
as variables.
|
||
|
||
- password (str): Encrypted or plaintext password.
|
||
|
||
- port (str): Port to connect to node, default 22 for ssh and 23
|
||
for telnet.
|
||
|
||
- protocol (str): Select ssh, telnet, kubectl or docker. Default is ssh.
|
||
|
||
- user (str): Username to of the node.
|
||
|
||
- config (obj): Pass the object created with class configfile with
|
||
key for decryption and extra configuration if you
|
||
are using connection manager.
|
||
|
||
- tags (dict) : Tags useful for automation and personal porpuse
|
||
like "os", "prompt" and "screenleght_command"
|
||
|
||
- jumphost (str): Reference another node to be used as a jumphost
|
||
</code></pre></div>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="connpy.node.interact"><code class="name flex">
|
||
<span>def <span class="ident">interact</span></span>(<span>self, debug=False, logger=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def interact(self, debug=False, logger=None):
|
||
'''
|
||
Asynchronous interactive session using Smart Tunnel architecture.
|
||
Allows multiplexing I/O and handling SIGWINCH events locally without blocking.
|
||
'''
|
||
connect = self._connect(debug=debug, logger=logger)
|
||
if connect == True:
|
||
try:
|
||
self._setup_interact_environment(debug=debug, logger=logger, async_mode=True)
|
||
|
||
local_stream = LocalStream()
|
||
|
||
def resize_callback(rows, cols):
|
||
try:
|
||
self.child.setwinsize(rows, cols)
|
||
except Exception:
|
||
pass
|
||
|
||
# Build local copilot handler
|
||
copilot_handler = self._build_local_copilot_handler()
|
||
|
||
asyncio.run(self._async_interact_loop(local_stream, resize_callback, copilot_handler=copilot_handler))
|
||
finally:
|
||
self._teardown_interact_environment()
|
||
else:
|
||
if logger:
|
||
logger("error", str(connect))
|
||
else:
|
||
printer.error(f"Connection failed: {str(connect)}")
|
||
sys.exit(1)</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Asynchronous interactive session using Smart Tunnel architecture.
|
||
Allows multiplexing I/O and handling SIGWINCH events locally without blocking.</p></div>
|
||
</dd>
|
||
<dt id="connpy.node.run"><code class="name flex">
|
||
<span>def <span class="ident">run</span></span>(<span>self,<br>commands,<br>vars=None,<br>*,<br>folder='',<br>prompt='>$|#$|\\$$|>.$|#.$|\\$.$',<br>stdout=False,<br>timeout=10,<br>logger=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def run(self, commands, vars = None,*, folder = '', prompt = r'>$|#$|\$$|>.$|#.$|\$.$', stdout = False, timeout = 10, logger = None):
|
||
'''
|
||
Run a command or list of commands on the node and return the output.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables
|
||
used in commands parameter.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave
|
||
empty to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- stdout (bool):Set True to send the command output to stdout.
|
||
default False.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
### Returns:
|
||
|
||
str: Output of the commands you ran on the node.
|
||
|
||
'''
|
||
connect = self._connect(timeout = timeout, logger = logger)
|
||
now = datetime.datetime.now().strftime('%Y-%m-%d_%H%M%S')
|
||
if connect == True:
|
||
if logger:
|
||
port_str = f":{self.port}" if self.port and self.protocol not in ["ssm", "kubectl", "docker"] else ""
|
||
logger("success", f"Connected to {self.unique} at {self.host}{port_str} via: {self.protocol}")
|
||
|
||
# Attempt to set the terminal size
|
||
try:
|
||
self.child.setwinsize(65535, 65535)
|
||
except Exception:
|
||
try:
|
||
self.child.setwinsize(10000, 10000)
|
||
except Exception:
|
||
pass
|
||
if "prompt" in self.tags:
|
||
prompt = self.tags["prompt"]
|
||
expects = [prompt, pexpect.EOF, pexpect.TIMEOUT]
|
||
output = ''
|
||
status = ''
|
||
if not isinstance(commands, list):
|
||
commands = [commands]
|
||
if "screen_length_command" in self.tags:
|
||
commands.insert(0, self.tags["screen_length_command"])
|
||
self.mylog = io.BytesIO()
|
||
self.child.logfile_read = self.mylog
|
||
for c in commands:
|
||
if vars is not None:
|
||
try:
|
||
c = c.format(**vars)
|
||
except KeyError as e:
|
||
self.output = f"Error: Variable {e} not defined in task or inventory"
|
||
self.status = 1
|
||
return self.output
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.sendline(c)
|
||
if result == 2:
|
||
break
|
||
if not result == 2:
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.close()
|
||
output = self._logclean(self.mylog.getvalue().decode(), True)
|
||
if logger:
|
||
logger("output", output)
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(output)
|
||
f.close()
|
||
self.output = output
|
||
if result == 2:
|
||
self.status = 2
|
||
else:
|
||
self.status = 0
|
||
return output
|
||
else:
|
||
self.output = connect
|
||
self.status = 1
|
||
if logger:
|
||
logger("error", f"Connection failed: {connect}")
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(connect)
|
||
|
||
f.close()
|
||
return connect</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Run a command or list of commands on the node and return the output.</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- vars (dict): Dictionary containing the definition of variables
|
||
used in commands parameter.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
</code></pre>
|
||
<h3 id="optional-named-parameters">Optional Named Parameters:</h3>
|
||
<pre><code>- folder (str): Path where output log should be stored, leave
|
||
empty to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- stdout (bool):Set True to send the command output to stdout.
|
||
default False.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>str: Output of the commands you ran on the node.
|
||
</code></pre></div>
|
||
</dd>
|
||
<dt id="connpy.node.test"><code class="name flex">
|
||
<span>def <span class="ident">test</span></span>(<span>self,<br>commands,<br>expected,<br>vars=None,<br>*,<br>folder='',<br>prompt='>$|#$|\\$$|>.$|#.$|\\$.$',<br>timeout=10,<br>logger=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def test(self, commands, expected, vars = None,*, folder = '', prompt = r'>$|#$|\$$|>.$|#.$|\$.$', timeout = 10, logger = None):
|
||
'''
|
||
Run a command or list of commands on the node, then check if expected value appears on the output after the last command.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
|
||
- expected (str) : Expected text to appear after running
|
||
all the commands on the node.You can use
|
||
variables as {varname} and defining them
|
||
in optional parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables
|
||
used in commands and expected parameters.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave
|
||
empty to not store logs.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
### Returns:
|
||
bool: true if expected value is found after running the commands
|
||
false if prompt is found before.
|
||
|
||
'''
|
||
now = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
|
||
connect = self._connect(timeout = timeout, logger = logger)
|
||
if connect == True:
|
||
if logger:
|
||
port_str = f":{self.port}" if self.port and self.protocol not in ["ssm", "kubectl", "docker"] else ""
|
||
logger("success", f"Connected to {self.unique} at {self.host}{port_str} via: {self.protocol}")
|
||
|
||
# Attempt to set the terminal size
|
||
try:
|
||
self.child.setwinsize(65535, 65535)
|
||
except Exception:
|
||
try:
|
||
self.child.setwinsize(10000, 10000)
|
||
except Exception:
|
||
pass
|
||
if "prompt" in self.tags:
|
||
prompt = self.tags["prompt"]
|
||
expects = [prompt, pexpect.EOF, pexpect.TIMEOUT]
|
||
output = ''
|
||
if not isinstance(commands, list):
|
||
commands = [commands]
|
||
if not isinstance(expected, list):
|
||
expected = [expected]
|
||
if "screen_length_command" in self.tags:
|
||
commands.insert(0, self.tags["screen_length_command"])
|
||
self.mylog = io.BytesIO()
|
||
self.child.logfile_read = self.mylog
|
||
for c in commands:
|
||
if vars is not None:
|
||
try:
|
||
c = c.format(**vars)
|
||
except KeyError as e:
|
||
self.output = f"Error: Variable {e} not defined in task or inventory"
|
||
self.status = 1
|
||
return self.output
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.sendline(c)
|
||
if result == 2:
|
||
break
|
||
if not result == 2:
|
||
result = self.child.expect(expects, timeout = timeout)
|
||
self.child.close()
|
||
output = self._logclean(self.mylog.getvalue().decode(), True)
|
||
if logger:
|
||
logger("output", output)
|
||
if folder != '':
|
||
with open(folder + "/" + self.unique + "_" + now + ".txt", "w") as f:
|
||
f.write(output)
|
||
f.close()
|
||
self.output = output
|
||
if result in [0, 1]:
|
||
# lastcommand = commands[-1]
|
||
# if vars is not None:
|
||
# lastcommand = lastcommand.format(**vars)
|
||
# last_command_index = output.rfind(lastcommand)
|
||
# cleaned_output = output[last_command_index + len(lastcommand):].strip()
|
||
self.result = {}
|
||
for e in expected:
|
||
if vars is not None:
|
||
e = e.format(**vars)
|
||
updatedprompt = re.sub(r'(?<!\\)\$', '', prompt)
|
||
newpattern = f".*({updatedprompt}).*{e}.*"
|
||
cleaned_output = output
|
||
cleaned_output = re.sub(newpattern, '', cleaned_output)
|
||
if e in cleaned_output:
|
||
self.result[e] = True
|
||
else:
|
||
self.result[e]= False
|
||
self.status = 0
|
||
return self.result
|
||
if result == 2:
|
||
self.result = None
|
||
self.status = 2
|
||
return output
|
||
else:
|
||
self.result = None
|
||
self.output = connect
|
||
self.status = 1
|
||
return connect</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Run a command or list of commands on the node, then check if expected value appears on the output after the last command.</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- commands (str/list): Commands to run on the node. Should be
|
||
str or a list of str. You can use variables
|
||
as {varname} and defining them in optional
|
||
parameter vars.
|
||
|
||
- expected (str) : Expected text to appear after running
|
||
all the commands on the node.You can use
|
||
variables as {varname} and defining them
|
||
in optional parameter vars.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- vars (dict): Dictionary containing the definition of variables
|
||
used in commands and expected parameters.
|
||
Keys: Variable names.
|
||
Values: strings.
|
||
</code></pre>
|
||
<h3 id="optional-named-parameters">Optional Named Parameters:</h3>
|
||
<pre><code>- folder (str): Path where output log should be stored, leave
|
||
empty to not store logs.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
- timeout (int):Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>bool: true if expected value is found after running the commands
|
||
false if prompt is found before.
|
||
</code></pre></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
<dt id="connpy.nodes"><code class="flex name class">
|
||
<span>class <span class="ident">nodes</span></span>
|
||
<span>(</span><span>nodes: dict, config='')</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@ClassHook
|
||
class nodes:
|
||
''' This class generates a nodes object. Contains a list of node class objects and methods to run multiple tasks on nodes simultaneously.
|
||
|
||
### Attributes:
|
||
|
||
- nodelist (list): List of node class objects passed to the init
|
||
function.
|
||
|
||
- output (dict): Dictionary formed by nodes unique as keys,
|
||
output of the commands you ran on the node as
|
||
value. Created after running methods run or test.
|
||
|
||
- result (dict): Dictionary formed by nodes unique as keys, value
|
||
is True if expected value is found after running
|
||
the commands, False if prompt is found before.
|
||
Created after running method test.
|
||
|
||
- status (dict): Dictionary formed by nodes unique as keys, value:
|
||
0 if method run or test ended successfully.
|
||
1 if connection failed.
|
||
2 if expect timeouts without prompt or EOF.
|
||
|
||
- <unique> (obj): For each item in nodelist, there is an attribute
|
||
generated with the node unique.
|
||
'''
|
||
|
||
def __init__(self, nodes: dict, config = ''):
|
||
'''
|
||
### Parameters:
|
||
|
||
- nodes (dict): Dictionary formed by node information:
|
||
Keys: Unique name for each node.
|
||
Mandatory Subkeys: host(str).
|
||
Optional Subkeys: options(str), logs(str), password(str),
|
||
port(str), protocol(str), user(str).
|
||
For reference on subkeys check node class.
|
||
|
||
### Optional Parameters:
|
||
|
||
- config (obj): Pass the object created with class configfile with key
|
||
for decryption and extra configuration if you are using
|
||
connection manager.
|
||
'''
|
||
self.nodelist = []
|
||
self.config = config
|
||
for n in nodes:
|
||
this = node(n, **nodes[n], config = config)
|
||
self.nodelist.append(this)
|
||
setattr(self,n,this)
|
||
|
||
|
||
@MethodHook
|
||
def _splitlist(self, lst, n):
|
||
#split a list in lists of n members.
|
||
for i in range(0, len(lst), n):
|
||
yield lst[i:i + n]
|
||
|
||
|
||
@MethodHook
|
||
def run(self, commands, vars = None,*, folder = None, prompt = None, stdout = None, parallel = 10, timeout = None, on_complete = None, logger = None):
|
||
'''
|
||
Run a command or list of commands on all the nodes in nodelist.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the nodes. Should be str or
|
||
list of str. You can use variables as {varname}
|
||
and defining them in optional parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands parameter.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave empty
|
||
to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while routers
|
||
use ">" or "#". The default value should work for
|
||
most nodes. Change it if your connection need some
|
||
special symbol.
|
||
|
||
- stdout (bool): Set True to send the command output to stdout.
|
||
Default False.
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
|
||
###Returns:
|
||
|
||
dict: Dictionary formed by nodes unique as keys, Output of the
|
||
commands you ran on the node as value.
|
||
|
||
'''
|
||
args = {}
|
||
nodesargs = {}
|
||
args["commands"] = commands
|
||
if folder != None:
|
||
args["folder"] = folder
|
||
Path(folder).mkdir(parents=True, exist_ok=True)
|
||
if prompt != None:
|
||
args["prompt"] = prompt
|
||
if stdout != None and on_complete is None:
|
||
args["stdout"] = stdout
|
||
if timeout != None:
|
||
args["timeout"] = timeout
|
||
output = {}
|
||
status = {}
|
||
tasks = []
|
||
|
||
def _run_node(node_obj, node_args, callback):
|
||
"""Wrapper that runs a node and fires the callback on completion."""
|
||
node_obj.run(**node_args)
|
||
if callback:
|
||
callback(node_obj.unique, node_obj.output, node_obj.status)
|
||
|
||
for n in self.nodelist:
|
||
nodesargs[n.unique] = deepcopy(args)
|
||
if vars != None:
|
||
nodesargs[n.unique]["vars"] = {}
|
||
if "__global__" in vars.keys():
|
||
nodesargs[n.unique]["vars"].update(vars["__global__"])
|
||
for var_key, var_val in vars.items():
|
||
if var_key == "__global__":
|
||
continue
|
||
try:
|
||
if re.search(var_key, n.unique, re.IGNORECASE):
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
except re.error:
|
||
if var_key == n.unique:
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
|
||
# Pass the logger to the node
|
||
nodesargs[n.unique]["logger"] = logger
|
||
|
||
if on_complete:
|
||
tasks.append(threading.Thread(target=_run_node, args=(n, nodesargs[n.unique], on_complete)))
|
||
else:
|
||
tasks.append(threading.Thread(target=n.run, kwargs=nodesargs[n.unique]))
|
||
|
||
taskslist = list(self._splitlist(tasks, parallel))
|
||
|
||
for t in taskslist:
|
||
for i in t:
|
||
i.start()
|
||
for i in t:
|
||
i.join()
|
||
for i in self.nodelist:
|
||
output[i.unique] = i.output
|
||
status[i.unique] = i.status
|
||
self.output = output
|
||
self.status = status
|
||
return output
|
||
|
||
@MethodHook
|
||
def test(self, commands, expected, vars = None,*, folder = None, prompt = None, parallel = 10, timeout = None, on_complete = None, logger = None):
|
||
'''
|
||
Run a command or list of commands on all the nodes in nodelist, then check if expected value appears on the output after the last command.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be str or
|
||
list of str.
|
||
|
||
- expected (str) : Expected text to appear after running all the
|
||
commands on the node.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands and expected parameters.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary formed by nodes unique as keys, value is True if
|
||
expected value is found after running the commands, False
|
||
if prompt is found before.
|
||
|
||
'''
|
||
args = {}
|
||
nodesargs = {}
|
||
args["commands"] = commands
|
||
args["expected"] = expected
|
||
if folder != None:
|
||
args["folder"] = folder
|
||
Path(folder).mkdir(parents=True, exist_ok=True)
|
||
if prompt != None:
|
||
args["prompt"] = prompt
|
||
if timeout != None:
|
||
args["timeout"] = timeout
|
||
output = {}
|
||
result = {}
|
||
status = {}
|
||
tasks = []
|
||
|
||
def _test_node(node_obj, node_args, callback):
|
||
"""Wrapper that runs a node test and fires the callback on completion."""
|
||
node_obj.test(**node_args)
|
||
if callback:
|
||
callback(node_obj.unique, node_obj.output, node_obj.status, node_obj.result)
|
||
|
||
for n in self.nodelist:
|
||
nodesargs[n.unique] = deepcopy(args)
|
||
if vars != None:
|
||
nodesargs[n.unique]["vars"] = {}
|
||
if "__global__" in vars.keys():
|
||
nodesargs[n.unique]["vars"].update(vars["__global__"])
|
||
for var_key, var_val in vars.items():
|
||
if var_key == "__global__":
|
||
continue
|
||
try:
|
||
if re.search(var_key, n.unique, re.IGNORECASE):
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
except re.error:
|
||
if var_key == n.unique:
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
nodesargs[n.unique]["logger"] = logger
|
||
|
||
if on_complete:
|
||
tasks.append(threading.Thread(target=_test_node, args=(n, nodesargs[n.unique], on_complete)))
|
||
else:
|
||
tasks.append(threading.Thread(target=n.test, kwargs=nodesargs[n.unique]))
|
||
|
||
taskslist = list(self._splitlist(tasks, parallel))
|
||
for t in taskslist:
|
||
for i in t:
|
||
i.start()
|
||
for i in t:
|
||
i.join()
|
||
for i in self.nodelist:
|
||
result[i.unique] = i.result
|
||
output[i.unique] = i.output
|
||
status[i.unique] = i.status
|
||
self.output = output
|
||
self.result = result
|
||
self.status = status
|
||
return result</code></pre>
|
||
</details>
|
||
<div class="desc"><p>This class generates a nodes object. Contains a list of node class objects and methods to run multiple tasks on nodes simultaneously.</p>
|
||
<h3 id="attributes">Attributes:</h3>
|
||
<pre><code>- nodelist (list): List of node class objects passed to the init
|
||
function.
|
||
|
||
- output (dict): Dictionary formed by nodes unique as keys,
|
||
output of the commands you ran on the node as
|
||
value. Created after running methods run or test.
|
||
|
||
- result (dict): Dictionary formed by nodes unique as keys, value
|
||
is True if expected value is found after running
|
||
the commands, False if prompt is found before.
|
||
Created after running method test.
|
||
|
||
- status (dict): Dictionary formed by nodes unique as keys, value:
|
||
0 if method run or test ended successfully.
|
||
1 if connection failed.
|
||
2 if expect timeouts without prompt or EOF.
|
||
|
||
- <unique> (obj): For each item in nodelist, there is an attribute
|
||
generated with the node unique.
|
||
</code></pre>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- nodes (dict): Dictionary formed by node information:
|
||
Keys: Unique name for each node.
|
||
Mandatory Subkeys: host(str).
|
||
Optional Subkeys: options(str), logs(str), password(str),
|
||
port(str), protocol(str), user(str).
|
||
For reference on subkeys check node class.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- config (obj): Pass the object created with class configfile with key
|
||
for decryption and extra configuration if you are using
|
||
connection manager.
|
||
</code></pre></div>
|
||
<h3>Methods</h3>
|
||
<dl>
|
||
<dt id="connpy.nodes.run"><code class="name flex">
|
||
<span>def <span class="ident">run</span></span>(<span>self,<br>commands,<br>vars=None,<br>*,<br>folder=None,<br>prompt=None,<br>stdout=None,<br>parallel=10,<br>timeout=None,<br>on_complete=None,<br>logger=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def run(self, commands, vars = None,*, folder = None, prompt = None, stdout = None, parallel = 10, timeout = None, on_complete = None, logger = None):
|
||
'''
|
||
Run a command or list of commands on all the nodes in nodelist.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the nodes. Should be str or
|
||
list of str. You can use variables as {varname}
|
||
and defining them in optional parameter vars.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands parameter.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- folder (str): Path where output log should be stored, leave empty
|
||
to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while routers
|
||
use ">" or "#". The default value should work for
|
||
most nodes. Change it if your connection need some
|
||
special symbol.
|
||
|
||
- stdout (bool): Set True to send the command output to stdout.
|
||
Default False.
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
|
||
###Returns:
|
||
|
||
dict: Dictionary formed by nodes unique as keys, Output of the
|
||
commands you ran on the node as value.
|
||
|
||
'''
|
||
args = {}
|
||
nodesargs = {}
|
||
args["commands"] = commands
|
||
if folder != None:
|
||
args["folder"] = folder
|
||
Path(folder).mkdir(parents=True, exist_ok=True)
|
||
if prompt != None:
|
||
args["prompt"] = prompt
|
||
if stdout != None and on_complete is None:
|
||
args["stdout"] = stdout
|
||
if timeout != None:
|
||
args["timeout"] = timeout
|
||
output = {}
|
||
status = {}
|
||
tasks = []
|
||
|
||
def _run_node(node_obj, node_args, callback):
|
||
"""Wrapper that runs a node and fires the callback on completion."""
|
||
node_obj.run(**node_args)
|
||
if callback:
|
||
callback(node_obj.unique, node_obj.output, node_obj.status)
|
||
|
||
for n in self.nodelist:
|
||
nodesargs[n.unique] = deepcopy(args)
|
||
if vars != None:
|
||
nodesargs[n.unique]["vars"] = {}
|
||
if "__global__" in vars.keys():
|
||
nodesargs[n.unique]["vars"].update(vars["__global__"])
|
||
for var_key, var_val in vars.items():
|
||
if var_key == "__global__":
|
||
continue
|
||
try:
|
||
if re.search(var_key, n.unique, re.IGNORECASE):
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
except re.error:
|
||
if var_key == n.unique:
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
|
||
# Pass the logger to the node
|
||
nodesargs[n.unique]["logger"] = logger
|
||
|
||
if on_complete:
|
||
tasks.append(threading.Thread(target=_run_node, args=(n, nodesargs[n.unique], on_complete)))
|
||
else:
|
||
tasks.append(threading.Thread(target=n.run, kwargs=nodesargs[n.unique]))
|
||
|
||
taskslist = list(self._splitlist(tasks, parallel))
|
||
|
||
for t in taskslist:
|
||
for i in t:
|
||
i.start()
|
||
for i in t:
|
||
i.join()
|
||
for i in self.nodelist:
|
||
output[i.unique] = i.output
|
||
status[i.unique] = i.status
|
||
self.output = output
|
||
self.status = status
|
||
return output</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Run a command or list of commands on all the nodes in nodelist.</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- commands (str/list): Commands to run on the nodes. Should be str or
|
||
list of str. You can use variables as {varname}
|
||
and defining them in optional parameter vars.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands parameter.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
</code></pre>
|
||
<h3 id="optional-named-parameters">Optional Named Parameters:</h3>
|
||
<pre><code>- folder (str): Path where output log should be stored, leave empty
|
||
to disable logging.
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while routers
|
||
use ">" or "#". The default value should work for
|
||
most nodes. Change it if your connection need some
|
||
special symbol.
|
||
|
||
- stdout (bool): Set True to send the command output to stdout.
|
||
Default False.
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>dict: Dictionary formed by nodes unique as keys, Output of the
|
||
commands you ran on the node as value.
|
||
</code></pre></div>
|
||
</dd>
|
||
<dt id="connpy.nodes.test"><code class="name flex">
|
||
<span>def <span class="ident">test</span></span>(<span>self,<br>commands,<br>expected,<br>vars=None,<br>*,<br>folder=None,<br>prompt=None,<br>parallel=10,<br>timeout=None,<br>on_complete=None,<br>logger=None)</span>
|
||
</code></dt>
|
||
<dd>
|
||
<details class="source">
|
||
<summary>
|
||
<span>Expand source code</span>
|
||
</summary>
|
||
<pre><code class="python">@MethodHook
|
||
def test(self, commands, expected, vars = None,*, folder = None, prompt = None, parallel = 10, timeout = None, on_complete = None, logger = None):
|
||
'''
|
||
Run a command or list of commands on all the nodes in nodelist, then check if expected value appears on the output after the last command.
|
||
|
||
|
||
### Parameters:
|
||
|
||
- commands (str/list): Commands to run on the node. Should be str or
|
||
list of str.
|
||
|
||
- expected (str) : Expected text to appear after running all the
|
||
commands on the node.
|
||
|
||
### Optional Parameters:
|
||
|
||
- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands and expected parameters.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
|
||
### Optional Named Parameters:
|
||
|
||
- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
|
||
### Returns:
|
||
|
||
dict: Dictionary formed by nodes unique as keys, value is True if
|
||
expected value is found after running the commands, False
|
||
if prompt is found before.
|
||
|
||
'''
|
||
args = {}
|
||
nodesargs = {}
|
||
args["commands"] = commands
|
||
args["expected"] = expected
|
||
if folder != None:
|
||
args["folder"] = folder
|
||
Path(folder).mkdir(parents=True, exist_ok=True)
|
||
if prompt != None:
|
||
args["prompt"] = prompt
|
||
if timeout != None:
|
||
args["timeout"] = timeout
|
||
output = {}
|
||
result = {}
|
||
status = {}
|
||
tasks = []
|
||
|
||
def _test_node(node_obj, node_args, callback):
|
||
"""Wrapper that runs a node test and fires the callback on completion."""
|
||
node_obj.test(**node_args)
|
||
if callback:
|
||
callback(node_obj.unique, node_obj.output, node_obj.status, node_obj.result)
|
||
|
||
for n in self.nodelist:
|
||
nodesargs[n.unique] = deepcopy(args)
|
||
if vars != None:
|
||
nodesargs[n.unique]["vars"] = {}
|
||
if "__global__" in vars.keys():
|
||
nodesargs[n.unique]["vars"].update(vars["__global__"])
|
||
for var_key, var_val in vars.items():
|
||
if var_key == "__global__":
|
||
continue
|
||
try:
|
||
if re.search(var_key, n.unique, re.IGNORECASE):
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
except re.error:
|
||
if var_key == n.unique:
|
||
nodesargs[n.unique]["vars"].update(var_val)
|
||
nodesargs[n.unique]["logger"] = logger
|
||
|
||
if on_complete:
|
||
tasks.append(threading.Thread(target=_test_node, args=(n, nodesargs[n.unique], on_complete)))
|
||
else:
|
||
tasks.append(threading.Thread(target=n.test, kwargs=nodesargs[n.unique]))
|
||
|
||
taskslist = list(self._splitlist(tasks, parallel))
|
||
for t in taskslist:
|
||
for i in t:
|
||
i.start()
|
||
for i in t:
|
||
i.join()
|
||
for i in self.nodelist:
|
||
result[i.unique] = i.result
|
||
output[i.unique] = i.output
|
||
status[i.unique] = i.status
|
||
self.output = output
|
||
self.result = result
|
||
self.status = status
|
||
return result</code></pre>
|
||
</details>
|
||
<div class="desc"><p>Run a command or list of commands on all the nodes in nodelist, then check if expected value appears on the output after the last command.</p>
|
||
<h3 id="parameters">Parameters:</h3>
|
||
<pre><code>- commands (str/list): Commands to run on the node. Should be str or
|
||
list of str.
|
||
|
||
- expected (str) : Expected text to appear after running all the
|
||
commands on the node.
|
||
</code></pre>
|
||
<h3 id="optional-parameters">Optional Parameters:</h3>
|
||
<pre><code>- vars (dict): Dictionary containing the definition of variables for
|
||
each node, used in commands and expected parameters.
|
||
Keys should be formed by nodes unique names. Use
|
||
special key name __global__ for global variables.
|
||
Subkeys: Variable names.
|
||
Values: strings.
|
||
</code></pre>
|
||
<h3 id="optional-named-parameters">Optional Named Parameters:</h3>
|
||
<pre><code>- prompt (str): Prompt to be expected after a command is finished
|
||
running. Usually linux uses ">" or EOF while
|
||
routers use ">" or "#". The default value should
|
||
work for most nodes. Change it if your connection
|
||
need some special symbol.
|
||
|
||
|
||
- parallel (int): Number of nodes to run the commands simultaneously.
|
||
Default is 10, if there are more nodes that this
|
||
value, nodes are groups in groups with max this
|
||
number of members.
|
||
|
||
- timeout (int): Time in seconds for expect to wait for prompt/EOF.
|
||
default 10.
|
||
|
||
- on_complete (callable): Optional callback called when each node
|
||
finishes. Receives (unique, output, status).
|
||
Called from the node's thread so it must
|
||
be thread-safe.
|
||
</code></pre>
|
||
<h3 id="returns">Returns:</h3>
|
||
<pre><code>dict: Dictionary formed by nodes unique as keys, value is True if
|
||
expected value is found after running the commands, False
|
||
if prompt is found before.
|
||
</code></pre></div>
|
||
</dd>
|
||
</dl>
|
||
</dd>
|
||
</dl>
|
||
</section>
|
||
</article>
|
||
<nav id="sidebar">
|
||
<div class="toc">
|
||
<ul>
|
||
<li><a href="#connpy">Connpy</a><ul>
|
||
<li><a href="#ai-copilot-new-in-v6">🤖 AI Copilot (New in v6)</a></li>
|
||
<li><a href="#core-features">Core Features</a></li>
|
||
<li><a href="#installation">Installation</a><ul>
|
||
<li><a href="#run-it-in-windowslinux-using-docker">Run it in Windows/Linux using Docker</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#privacy-integration">🔒 Privacy & Integration</a><ul>
|
||
<li><a href="#privacy-policy">Privacy Policy</a></li>
|
||
<li><a href="#google-integration">Google Integration</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#usage">Usage</a><ul>
|
||
<li><a href="#basic-examples">Basic Examples:</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#plugin-requirements-for-connpy">Plugin Requirements for Connpy</a><ul>
|
||
<li><a href="#remote-plugin-execution">Remote Plugin Execution</a></li>
|
||
<li><a href="#general-structure">General Structure</a></li>
|
||
<li><a href="#preload-modifications-and-hooks">Preload Modifications and Hooks</a></li>
|
||
<li><a href="#command-completion-support">Command Completion Support</a></li>
|
||
</ul>
|
||
</li>
|
||
<li><a href="#grpc-service-architecture">⚙️ gRPC Service Architecture</a></li>
|
||
<li><a href="#automation-module-api">🐍 Automation Module (API)</a><ul>
|
||
<li><a href="#basic-execution">Basic Execution</a></li>
|
||
<li><a href="#parallel-tasks-with-variables">Parallel Tasks with Variables</a></li>
|
||
<li><a href="#ai-programmatic-use">AI Programmatic Use</a></li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</div>
|
||
<ul id="index">
|
||
<li><h3><a href="#header-submodules">Sub-modules</a></h3>
|
||
<ul>
|
||
<li><code><a title="connpy.cli" href="cli/index.html">connpy.cli</a></code></li>
|
||
<li><code><a title="connpy.grpc_layer" href="grpc_layer/index.html">connpy.grpc_layer</a></code></li>
|
||
<li><code><a title="connpy.mcp_client" href="mcp_client.html">connpy.mcp_client</a></code></li>
|
||
<li><code><a title="connpy.proto" href="proto/index.html">connpy.proto</a></code></li>
|
||
<li><code><a title="connpy.services" href="services/index.html">connpy.services</a></code></li>
|
||
<li><code><a title="connpy.tests" href="tests/index.html">connpy.tests</a></code></li>
|
||
<li><code><a title="connpy.tunnels" href="tunnels.html">connpy.tunnels</a></code></li>
|
||
<li><code><a title="connpy.utils" href="utils.html">connpy.utils</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li><h3><a href="#header-classes">Classes</a></h3>
|
||
<ul>
|
||
<li>
|
||
<h4><code><a title="connpy.Plugins" href="#connpy.Plugins">Plugins</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="connpy.Plugins.verify_script" href="#connpy.Plugins.verify_script">verify_script</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li>
|
||
<h4><code><a title="connpy.ai" href="#connpy.ai">ai</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="connpy.ai.SAFE_COMMANDS" href="#connpy.ai.SAFE_COMMANDS">SAFE_COMMANDS</a></code></li>
|
||
<li><code><a title="connpy.ai.aask_copilot" href="#connpy.ai.aask_copilot">aask_copilot</a></code></li>
|
||
<li><code><a title="connpy.ai.architect_system_prompt" href="#connpy.ai.architect_system_prompt">architect_system_prompt</a></code></li>
|
||
<li><code><a title="connpy.ai.ask" href="#connpy.ai.ask">ask</a></code></li>
|
||
<li><code><a title="connpy.ai.confirm" href="#connpy.ai.confirm">confirm</a></code></li>
|
||
<li><code><a title="connpy.ai.delete_session" href="#connpy.ai.delete_session">delete_session</a></code></li>
|
||
<li><code><a title="connpy.ai.engineer_system_prompt" href="#connpy.ai.engineer_system_prompt">engineer_system_prompt</a></code></li>
|
||
<li><code><a title="connpy.ai.get_last_session_id" href="#connpy.ai.get_last_session_id">get_last_session_id</a></code></li>
|
||
<li><code><a title="connpy.ai.get_node_info_tool" href="#connpy.ai.get_node_info_tool">get_node_info_tool</a></code></li>
|
||
<li><code><a title="connpy.ai.list_nodes_tool" href="#connpy.ai.list_nodes_tool">list_nodes_tool</a></code></li>
|
||
<li><code><a title="connpy.ai.list_sessions" href="#connpy.ai.list_sessions">list_sessions</a></code></li>
|
||
<li><code><a title="connpy.ai.load_session_data" href="#connpy.ai.load_session_data">load_session_data</a></code></li>
|
||
<li><code><a title="connpy.ai.manage_memory_tool" href="#connpy.ai.manage_memory_tool">manage_memory_tool</a></code></li>
|
||
<li><code><a title="connpy.ai.register_ai_tool" href="#connpy.ai.register_ai_tool">register_ai_tool</a></code></li>
|
||
<li><code><a title="connpy.ai.run_commands_tool" href="#connpy.ai.run_commands_tool">run_commands_tool</a></code></li>
|
||
<li><code><a title="connpy.ai.save_session" href="#connpy.ai.save_session">save_session</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li>
|
||
<h4><code><a title="connpy.configfile" href="#connpy.configfile">configfile</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="connpy.configfile.encrypt" href="#connpy.configfile.encrypt">encrypt</a></code></li>
|
||
<li><code><a title="connpy.configfile.getitem" href="#connpy.configfile.getitem">getitem</a></code></li>
|
||
<li><code><a title="connpy.configfile.getitems" href="#connpy.configfile.getitems">getitems</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li>
|
||
<h4><code><a title="connpy.node" href="#connpy.node">node</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="connpy.node.interact" href="#connpy.node.interact">interact</a></code></li>
|
||
<li><code><a title="connpy.node.run" href="#connpy.node.run">run</a></code></li>
|
||
<li><code><a title="connpy.node.test" href="#connpy.node.test">test</a></code></li>
|
||
</ul>
|
||
</li>
|
||
<li>
|
||
<h4><code><a title="connpy.nodes" href="#connpy.nodes">nodes</a></code></h4>
|
||
<ul class="">
|
||
<li><code><a title="connpy.nodes.run" href="#connpy.nodes.run">run</a></code></li>
|
||
<li><code><a title="connpy.nodes.test" href="#connpy.nodes.test">test</a></code></li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</li>
|
||
</ul>
|
||
</nav>
|
||
</main>
|
||
<footer id="footer">
|
||
<p>Generated by <a href="https://pdoc3.github.io/pdoc" title="pdoc: Python API documentation generator"><cite>pdoc</cite> 0.11.5</a>.</p>
|
||
</footer>
|
||
</body>
|
||
</html>
|