commit c8224bfb943d1f863dfb1386dc1edc21e168b507
parent e7830972547b44a7a1e6d265721ecf6a9fae0339
Author: Asher Morgan <asher@ashermorgan.net>
Date: Sat, 15 Aug 2026 15:22:08 -0700
Improve error connection handling
Diffstat:
2 files changed, 79 insertions(+), 40 deletions(-)
diff --git a/connection_plugins/http.py b/connection_plugins/http.py
@@ -18,9 +18,10 @@ DOCUMENTATION = r'''
# ruff: disable[E402]
import base64
import json
+import os
import typing as t
-from ansible.module_utils.common.text.converters import to_text
+from ansible.errors import AnsibleFileNotFound, AnsibleError
from ansible.plugins.connection import ConnectionBase
from ansible.utils.display import Display
# ruff: enable[E402]
@@ -58,7 +59,7 @@ class Connection(ConnectionBase):
def exec_command(self, cmd: str, in_data: bytes | None = None,
sudoable: bool = True) -> tuple[int, bytes, bytes]:
- """Run a command on the host (TODO)."""
+ """Run a command on the host."""
super(Connection, self).exec_command(cmd, in_data=in_data, sudoable=sudoable)
@@ -71,11 +72,13 @@ class Connection(ConnectionBase):
# assert sudoable is False
self.sendbuf.write(json.dumps({
- 'command': cmd,
+ 'exec': cmd,
}) + '\n')
self.sendbuf.flush()
res = json.loads(self.recvbuf.readline())
+ if 'err' in res:
+ raise AnsibleError(res['err'])
return (
res['returncode'],
res['stdout'].encode('latin1'),
@@ -84,7 +87,7 @@ class Connection(ConnectionBase):
def put_file(self, in_path: str, out_path: str) -> None:
- """Transfer file to host (TODO)."""
+ """Transfer file to host."""
super(Connection, self).put_file(in_path, out_path)
@@ -94,6 +97,8 @@ class Connection(ConnectionBase):
assert self.sendbuf is not None
assert self.recvbuf is not None
+ if not os.path.exists(in_path):
+ raise AnsibleFileNotFound(f'file or module does not exist: {in_path}')
with open(in_path, 'rb') as f:
data = base64.b64encode(f.read()).decode()
@@ -105,11 +110,13 @@ class Connection(ConnectionBase):
}) + '\n')
self.sendbuf.flush()
- self.recvbuf.readline()
+ res = json.loads(self.recvbuf.readline())
+ if 'err' in res:
+ raise AnsibleError(res['err'])
def fetch_file(self, in_path: str, out_path: str) -> None:
- """Fetch file from host (TODO)."""
+ """Fetch file from host."""
super(Connection, self).fetch_file(in_path, out_path)
@@ -120,11 +127,13 @@ class Connection(ConnectionBase):
assert self.recvbuf is not None
self.sendbuf.write(json.dumps({
- 'get': in_path,
+ 'fetch': in_path,
}) + '\n')
self.sendbuf.flush()
res = json.loads(self.recvbuf.readline())
+ if 'err' in res:
+ raise AnsibleError(res['err'])
with open(out_path, 'rb') as f:
f.write(base64.b64decode(res['data']))
diff --git a/install.py b/install.py
@@ -6,49 +6,79 @@ import time
import requests # TODO: eliminate dependency?
-API = 'http://localhost:5000'
-
-res = requests.post(f'{API}/runners/', json={
- 'host': 'localhost',
-})
-cookies = res.cookies
-assert res.status_code == 201
+API = 'http://localhost:5000'
-job_url = f'{API}{res.headers['Location']}'
-req = {}
-while True:
- res = requests.put(job_url, cookies=cookies, json=req).json()
- if res['status'] not in ['started', 'running']:
- break
+def exec(args):
+ print(f'EXEC {args}')
- req = {}
- if (command := res.get('command')) is not None:
- print('EXEC', command)
+ try:
p = subprocess.Popen(
- command,
+ args,
shell=True,
stdin=subprocess.PIPE,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
)
stdout, stderr = p.communicate()
- req['returncode'] = p.returncode
- req['stdout'] = stdout.decode('latin1')
- req['stderr'] = stderr.decode('latin1')
- print('EXEC RETURN', p.returncode)
-
- elif (put := res.get('put')) is not None:
- print('PUT', put['dest'], len(put['data']), 'bytes')
- with open(put['dest'], 'wb') as f:
- f.write(base64.b64decode(put['data']))
- req['ok'] = True
-
- elif (get := res.get('get')) is not None:
- print('PUT', put['dest'], len(put['data']), 'bytes')
- with open(put['dest'], 'rb') as f:
- req['data'] = base64.b64encode(f.read()).decode()
+ return {
+ 'returncode': p.returncode,
+ 'stdout': stdout.decode('latin1'),
+ 'stderr': stderr.decode('latin1'),
+ }
+ except Exception as e:
+ return { 'err': str(e) }
+
+def put(args):
+ print(f"PUT {args['dest']} ({len(args['data'])} bytes)")
+
+ try:
+ with open(args['dest'], 'wb') as f:
+ f.write(base64.b64decode(args['data']))
+ except Exception as e:
+ return { 'err': str(e) }
else:
- time.sleep(0.1)
+ return { 'ok': True }
+
+
+def fetch(args):
+ print(f"PUT {args['dest']} ({len(args['data'])} bytes)")
+
+ try:
+ with open(args['dest'], 'rb') as f:
+ return {
+ 'data': base64.b64encode(f.read()).decode(),
+ }
+ except Exception as e:
+ return { 'err': str(e) }
+
+
+def main():
+ res = requests.post(f'{API}/runners/', json={
+ 'host': 'localhost',
+ })
+ assert res.status_code == 201
+ cookies = res.cookies
+ job_url = f'{API}{res.headers['Location']}'
+
+ req = {}
+ while True:
+ res = requests.put(job_url, cookies=cookies, json=req).json()
+
+ req = {}
+ if res['status'] not in ['started', 'running']:
+ break
+ elif 'exec' in res:
+ req = exec(res['exec'])
+ elif 'put' in res:
+ req = put(res['put'])
+ elif 'fetch' in res:
+ req = fetch(res['fetch'])
+ else:
+ time.sleep(0.1)
+
+
+if __name__ == '__main__':
+ main()
