ominfra 0.0.0.dev7__py3-none-any.whl

ominfra/pyremote/bootstrap.py

@@ -0,0 +1,149 @@
+"""
+Basically this: https://mitogen.networkgenomics.com/howitworks.html
+"""
+import base64
+import inspect
+import os
+import sys
+import textwrap
+import typing as ta
+import zlib
+
+
+##
+
+
+_BOOTSTRAP_COMM_FD = 100
+_BOOTSTRAP_SRC_FD = 101
+
+_BOOTSTRAP_CHILD_PID_VAR = '_PYR_CPID'
+_BOOTSTRAP_ARGV0_VAR = '_PYR_ARGV0'
+
+BOOTSTRAP_ACK0 = b'OPYR000\n'
+BOOTSTRAP_ACK1 = b'OPYR001\n'
+
+_BOOTSTRAP_PROC_TITLE_FMT = '(pyremote:%s)'
+
+_BOOTSTRAP_IMPORTS = [
+    'base64',
+    'os',
+    'sys',
+    'zlib',
+]
+
+
+def _bootstrap_main(context_name: str, main_z_len: int) -> None:
+    # Two copies of main src to be sent to parent
+    r0, w0 = os.pipe()
+    r1, w1 = os.pipe()
+
+    if (cp := os.fork()):
+        # Parent process
+
+        # Dup original stdin to comm_fd for use as comm channel
+        os.dup2(0, _BOOTSTRAP_COMM_FD)
+
+        # Overwrite stdin (fed to python repl) with first copy of src
+        os.dup2(r0, 0)
+
+        # Dup second copy of src to src_fd to recover after launch
+        os.dup2(r1, _BOOTSTRAP_SRC_FD)
+
+        # Close remaining fd's
+        for f in [r0, w0, r1, w1]:
+            os.close(f)
+
+        # Save child pid to close after relaunch
+        os.environ[_BOOTSTRAP_CHILD_PID_VAR] = str(cp)
+
+        # Save original argv0
+        os.environ[_BOOTSTRAP_ARGV0_VAR] = sys.executable
+
+        # Start repl reading stdin from r0
+        os.execl(sys.executable, sys.executable + (_BOOTSTRAP_PROC_TITLE_FMT % (context_name,)))
+
+    else:
+        # Child process
+
+        # Write first ack
+        os.write(1, BOOTSTRAP_ACK0)
+
+        # Read main src from stdin
+        main_src = zlib.decompress(os.fdopen(0, 'rb').read(main_z_len))
+
+        # Write both copies of main src
+        for w in [w0, w1]:
+            fp = os.fdopen(w, 'wb', 0)
+            fp.write(main_src)
+            fp.close()
+
+        # Write second ack
+        os.write(1, BOOTSTRAP_ACK1)
+
+        sys.exit(0)
+
+
+#
+
+
+def bootstrap_payload(context_name: str, main_z_len: int) -> str:
+    bs_src = textwrap.dedent(inspect.getsource(_bootstrap_main))
+
+    for gl in [
+        '_BOOTSTRAP_COMM_FD',
+        '_BOOTSTRAP_SRC_FD',
+
+        '_BOOTSTRAP_CHILD_PID_VAR',
+        '_BOOTSTRAP_ARGV0_VAR',
+
+        'BOOTSTRAP_ACK0',
+        'BOOTSTRAP_ACK1',
+
+        '_BOOTSTRAP_PROC_TITLE_FMT',
+    ]:
+        bs_src = bs_src.replace(gl, repr(globals()[gl]))
+
+    bs_src = '\n'.join(
+        cl
+        for l in bs_src.splitlines()
+        if (cl := (l.split('#')[0]).rstrip())
+        if cl.strip()
+    )
+
+    bs_z = zlib.compress(bs_src.encode('utf-8'))
+    bs_z64 = base64.encodebytes(bs_z).replace(b'\n', b'')
+
+    stmts = [
+        f'import {", ".join(_BOOTSTRAP_IMPORTS)}',
+        f'exec(zlib.decompress(base64.decodebytes({bs_z64!r})))',
+        f'_bootstrap_main({context_name!r}, {main_z_len})',
+    ]
+
+    cmd = '; '.join(stmts)
+    return cmd
+
+
+#
+
+
+class PostBoostrap(ta.NamedTuple):
+    input: ta.BinaryIO
+    main_src: str
+
+
+def post_boostrap() -> PostBoostrap:
+    # Restore original argv0
+    sys.executable = os.environ.pop(_BOOTSTRAP_ARGV0_VAR)
+
+    # Reap boostrap child
+    os.waitpid(int(os.environ.pop(_BOOTSTRAP_CHILD_PID_VAR)), 0)
+
+    # Read second copy of main src
+    r1 = os.fdopen(_BOOTSTRAP_SRC_FD, 'rb', 0)
+    main_src = r1.read().decode('utf-8')
+    r1.close()
+
+    return PostBoostrap(
+        input=os.fdopen(_BOOTSTRAP_COMM_FD, 'rb', 0),
+        main_src=main_src,
+    )

ominfra/pyremote/runcommands.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+# @omdev-amalg ./_runcommands.py
+# ruff: noqa: UP006 UP007
+import dataclasses as dc
+import io
+import json
+import subprocess
+import sys
+import typing as ta
+
+from omlish.lite.json import json_dumps_compact
+from omlish.lite.marshal import marshal_obj
+from omlish.lite.marshal import unmarshal_obj
+from omlish.lite.subprocesses import subprocess_maybe_shell_wrap_exec
+
+from .bootstrap import post_boostrap
+
+
+@dc.dataclass(frozen=True)
+class CommandRequest:
+    cmd: ta.Sequence[str]
+    in_: ta.Optional[bytes] = None
+
+
+@dc.dataclass(frozen=True)
+class CommandResponse:
+    req: CommandRequest
+    rc: int
+    out: bytes
+    err: bytes
+
+
+def _run_commands_loop(input: ta.BinaryIO, output: ta.BinaryIO = sys.stdout.buffer) -> None:  # noqa
+    while (l := input.readline().decode('utf-8').strip()):
+        req: CommandRequest = unmarshal_obj(json.loads(l), CommandRequest)
+        proc = subprocess.Popen(  # type: ignore
+            subprocess_maybe_shell_wrap_exec(*req.cmd),
+            **(dict(stdin=io.BytesIO(req.in_)) if req.in_ is not None else {}),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+        )
+        out, err = proc.communicate()
+        resp = CommandResponse(
+            req=req,
+            rc=proc.returncode,
+            out=out,  # noqa
+            err=err,  # noqa
+        )
+        output.write(json_dumps_compact(marshal_obj(resp)).encode('utf-8'))
+        output.write(b'\n')
+        output.flush()
+
+
+def run_commands_main() -> None:
+    bs = post_boostrap()
+    _run_commands_loop(bs.input)

ominfra/ssh.py

@@ -0,0 +1,191 @@
+"""
+TODO:
+ - sessionized
+ - streamed
+ - actual timeout
+
+bcrypt
+fido2
+gssapi
+libnacl
+pkcs11
+pyOpenSSL
+
+asyncssh[bcrypt,fido2,gssapi,libnacl,pkcs11,pyOpenSSL]
+"""
+import asyncio
+import contextlib
+import shlex
+import typing as ta
+
+from omlish import check
+from omlish import dataclasses as dc
+from omlish import lang
+from omserv.secrets import load_secrets
+
+from .cmds import CommandRunner
+from .cmds import LocalCommandRunner
+
+
+if ta.TYPE_CHECKING:
+    import asyncssh
+    import paramiko
+else:
+    asyncssh = lang.proxy_import('asyncssh')
+    paramiko = lang.proxy_import('paramiko')
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class SshConfig:
+    host: str | None = None
+    port: int | None = None
+
+    username: str | None = None
+    password: str | None = None
+
+    key_file_path: str | None = None
+
+
+class SshSubprocessCommandRunner(CommandRunner):
+    def __init__(
+            self,
+            cfg: SshConfig,
+            lcr: LocalCommandRunner | None = None,
+    ) -> None:
+        super().__init__()
+        self._cfg = check.isinstance(cfg, SshConfig)
+        self._lcr = check.isinstance(lcr, LocalCommandRunner) if lcr is not None else LocalCommandRunner()
+
+    async def run_command(self, cmd: CommandRunner.Command) -> CommandRunner.Result:
+        args = ['ssh']
+
+        if self._cfg.key_file_path is not None:
+            args.extend(['-i', self._cfg.key_file_path])
+
+        if self._cfg.port is not None:
+            args.extend(['-p', str(self._cfg.port)])
+
+        dst = check.non_empty_str(self._cfg.host)
+        if self._cfg.username is not None:
+            dst = f'{self._cfg.username}@{dst}'
+        if self._cfg.password is not None:
+            raise NotImplementedError
+
+        args.append(dst)
+        args.extend(cmd.args)
+        lcmd = CommandRunner.Command(
+            args=args,
+            in_=cmd.in_,
+        )
+        return await self._lcr.run_command(lcmd)
+
+
+class AsyncsshSshCommandRunner(CommandRunner):
+    def __init__(
+            self,
+            cfg: SshConfig,
+    ) -> None:
+        super().__init__()
+        self._cfg = check.isinstance(cfg, SshConfig)
+
+    async def run_command(self, cmd: CommandRunner.Command) -> CommandRunner.Result:
+        arg = ' '.join(map(shlex.quote, cmd.args))
+
+        async with asyncssh.connect(
+                self._cfg.host,
+                encoding=None,
+                **(dict(port=int(self._cfg.port)) if self._cfg.port is not None else {}),
+                **(dict(username=self._cfg.username) if self._cfg.username is not None else {}),
+                **(dict(password=self._cfg.password) if self._cfg.password is not None else {}),
+                **(dict(client_keys=[self._cfg.key_file_path]) if self._cfg.key_file_path is not None else {}),
+                known_hosts=None,
+        ) as conn:
+            proc: asyncssh.SSHClientProcess
+            async with await conn.create_process(arg) as proc:
+                checkw = False
+                timeout = None
+                res = await proc.wait(checkw, timeout)
+
+        return CommandRunner.Result(
+            rc=check.not_none(res.returncode),
+            out=check.isinstance(res.stdout, bytes),
+            err=check.isinstance(res.stderr, bytes),
+        )
+
+
+class ParamikoSshCommandRunner(CommandRunner):
+    def __init__(
+            self,
+            cfg: SshConfig,
+    ) -> None:
+        super().__init__()
+        self._cfg = check.isinstance(cfg, SshConfig)
+
+    def _run_command(self, cmd: CommandRunner.Command) -> CommandRunner.Result:
+        arg = ' '.join(map(shlex.quote, cmd.args))
+
+        kw: dict[str, ta.Any] = {}
+        if self._cfg.port is not None:
+            kw.update(port=int(self._cfg.port))
+        if self._cfg.username is not None:
+            kw.update(username=self._cfg.username)
+        if self._cfg.password is not None:
+            kw.update(password=self._cfg.password)
+        if self._cfg.key_file_path is not None:
+            kw.update(key_filename=self._cfg.key_file_path)
+
+        client: paramiko.client.SSHClient
+        with contextlib.closing(paramiko.client.SSHClient()) as client:
+            client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+            client.connect(
+                check.not_none(self._cfg.host),
+                **kw,
+            )
+
+            si, so, se = client.exec_command(arg)
+            # https://stackoverflow.com/questions/60037299/attributeerror-nonetype-object-has-no-attribute-time-paramiko
+            si.close()
+
+            rc = so.channel.recv_exit_status()
+            out = so.read()
+            err = se.read()
+
+        return CommandRunner.Result(
+            rc=rc,
+            out=out,
+            err=err,
+        )
+
+    async def run_command(self, cmd: CommandRunner.Command) -> CommandRunner.Result:
+        return await asyncio.to_thread(self._run_command, cmd)
+
+
+async def _a_main() -> None:
+    cmd = CommandRunner.Command(
+        ['ls', '-al'],
+    )
+
+    cfg = load_secrets()
+
+    sc = SshConfig(
+        host=cfg['ec2_ssh_host'],
+        username=cfg['ec2_ssh_user'],
+        key_file_path=cfg['ec2_ssh_key_file'],
+    )
+
+    for scr in [
+        SshSubprocessCommandRunner(sc),
+        AsyncsshSshCommandRunner(sc),
+        ParamikoSshCommandRunner(sc),
+    ]:
+        rc = await scr.run_command(cmd)
+        check.equal(rc.rc, 0)
+        print(rc.out.decode())
+
+
+if __name__ == '__main__':
+    asyncio.run(_a_main())

ominfra/tools/listresources.py

@@ -0,0 +1,256 @@
+"""
+TODO:
+ - async
+ - 'things' not just servers ('resources'?) - s3 buckets, db servers, etc
+ - ssh
+  - keys
+
+unique server ids:
+ aws:{region_name}:{instance_id}
+ runpod:{id}
+ lambda_labs:{id}
+"""
+import json
+import pprint
+import typing as ta
+
+import urllib3
+
+from omlish import check
+from omlish import dataclasses as dc
+from omlish import lang
+from omserv.secrets import load_secrets
+
+
+##
+
+
+@lang.cached_function
+def _get_secrets() -> dict[str, ta.Any]:
+    return load_secrets()
+
+
+##
+
+
+class Resource(lang.Abstract):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class Server(Resource):
+    host: str
+
+
+class ObjectStorage(Resource):
+    pass
+
+
+@dc.dataclass(frozen=True)
+class DbInstance(Resource):
+    host: str
+    port: int
+
+
+##
+
+
+class AwsResource(Resource):
+    pass
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class Ec2Server(Server, AwsResource):
+    id: str
+    region: str
+
+
+def get_ec2_servers() -> list[Ec2Server]:
+    cfg = _get_secrets()
+    import boto3
+    session = boto3.Session(
+        aws_access_key_id=cfg['aws_access_key_id'],
+        aws_secret_access_key=cfg['aws_secret_access_key'],
+        region_name=cfg['aws_region'],
+    )
+    ec2 = session.client('ec2')
+    resp = ec2.describe_instances()
+    out: list[Ec2Server] = []
+    for res in resp.get('Reservations', []):
+        for inst in res.get('Instances', []):
+            out.append(Ec2Server(
+                host=inst['PublicIpAddress'],
+                id=inst['InstanceId'],
+                region=ec2.meta.region_name,
+            ))
+    return out
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class RdsInstance(DbInstance, AwsResource):
+    id: str
+    region: str
+    engine: str
+
+
+def get_rds_instances() -> list[RdsInstance]:
+    cfg = _get_secrets()
+    import boto3
+    session = boto3.Session(
+        aws_access_key_id=cfg['aws_access_key_id'],
+        aws_secret_access_key=cfg['aws_secret_access_key'],
+        region_name=cfg['aws_region'],
+    )
+    rds = session.client('rds')
+    resp = rds.describe_db_instances()
+    out: list[RdsInstance] = []
+    for inst in resp.get('DBInstances', []):
+        out.append(RdsInstance(
+            host=inst['Endpoint']['Address'],
+            port=inst['Endpoint']['Port'],
+            id=inst['DBInstanceIdentifier'],
+            region=rds.meta.region_name,
+            engine=inst['Engine'],
+        ))
+    return out
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class GcpServer(Server):
+    id: str
+    zone: str
+
+
+def get_gcp_servers() -> list[GcpServer]:
+    cfg = _get_secrets()
+    from google.oauth2 import service_account
+    credentials = service_account.Credentials.from_service_account_info(cfg['gcp_oauth2'])
+    from google.cloud import compute_v1
+    instance_client = compute_v1.InstancesClient(credentials=credentials)
+    request = compute_v1.AggregatedListInstancesRequest()
+    request.project = cfg['gcp_project_id']
+    request.max_results = 50
+    out: list[GcpServer] = []
+    for zone, response in instance_client.aggregated_list(request=request):
+        for instance in (response.instances or []):
+            ip = check.single([ac.nat_i_p for ni in instance.network_interfaces for ac in ni.access_configs if ac.nat_i_p])  # noqa
+            out.append(GcpServer(
+                host=ip,
+                id=instance.name,
+                zone=zone,
+            ))
+    return out
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class RunpodServer(Server):
+    id: str
+
+
+def get_runpod_servers() -> list[RunpodServer]:
+    api_key = _get_secrets()['runpod_api_key']
+    query = 'query Pods { myself { pods { id runtime { ports { ip isIpPublic privatePort publicPort type } } } } }'
+    resp = urllib3.request(
+        'POST',
+        f'https://api.runpod.io/graphql?api_key={api_key}',
+        body=('{"query": "' + query + '"}').encode('utf-8'),
+        headers={
+            'content-type': 'application/json',
+        },
+    )
+    dct = json.loads(resp.data.decode('utf-8')).get('data', {}).get('myself', {})
+    out = []
+    for pod in dct.get('pods', []):
+        ssh = check.single([p for p in pod['runtime']['ports'] if p['isIpPublic'] and p['privatePort'] == 22])
+        out.append(RunpodServer(
+            host=f'{ssh["ip"]}:{ssh["publicPort"]}',
+            id=pod['id'],
+        ))
+    return out
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class LambdaLabsServer(Server):
+    id: str
+
+
+def get_lambda_labs_servers() -> list[LambdaLabsServer]:
+    api_key = _get_secrets()['lambda_labs_api_key']
+    resp = urllib3.request(
+        'GET',
+        'https://cloud.lambdalabs.com/api/v1/instances',
+        headers=urllib3.make_headers(
+            basic_auth=f'{api_key}:',
+        ),
+    )
+    insts = json.loads(resp.data.decode('utf-8')).get('data', {})
+    out: list[LambdaLabsServer] = []
+    for inst in insts:
+        out.append(LambdaLabsServer(
+            host=inst['ip'],
+            id=inst['id'],
+        ))
+    return out
+
+
+##
+
+
+@dc.dataclass(frozen=True)
+class DigitalOceanServer(Server):
+    id: str
+
+
+def get_digital_ocean_servers() -> list[DigitalOceanServer]:
+    api_key = _get_secrets()['digital_ocean_api_key']
+    resp = urllib3.request(
+        'GET',
+        'https://api.digitalocean.com/v2/droplets',
+        headers={
+            'Authorization': f'Bearer {api_key}',
+        },
+    )
+    droplets = json.loads(resp.data.decode('utf-8')).get('droplets', [])
+    out: list[DigitalOceanServer] = []
+    for droplet in droplets:
+        net = check.single([n for n in droplet['networks']['v4'] if n['type'] == 'public'])
+        out.append(DigitalOceanServer(
+            host=net['ip_address'],
+            id=droplet['id'],
+        ))
+    return out
+
+
+##
+
+
+def _main() -> None:
+    rsrcs: list[Resource] = [
+        *get_ec2_servers(),
+        *get_rds_instances(),
+        *get_gcp_servers(),
+        *get_runpod_servers(),
+        *get_lambda_labs_servers(),
+        *get_digital_ocean_servers(),
+    ]
+
+    pprint.pprint(rsrcs)
+
+
+if __name__ == '__main__':
+    _main()

ominfra-0.0.0.dev7.dist-info/LICENSE

@@ -0,0 +1,21 @@
+Copyright 2023- wrmsr
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
+following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following
+disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following
+disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products
+derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
+INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
+SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
+WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
+THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

ominfra-0.0.0.dev7.dist-info/METADATA

@@ -0,0 +1,19 @@
+Metadata-Version: 2.1
+Name: ominfra
+Version: 0.0.0.dev7
+Summary: ominfra
+Author: wrmsr
+License: BSD-3-Clause
+Project-URL: source, https://github.com/wrmsr/omlish
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: OS Independent
+Classifier: Operating System :: POSIX
+Requires-Python: >=3.12
+License-File: LICENSE
+Requires-Dist: omlish ==0.0.0.dev7
+Provides-Extra: ssh
+Requires-Dist: paramiko >=3.4 ; extra == 'ssh'
+Requires-Dist: asyncssh >=2.16 ; (python_version < "3.13") and extra == 'ssh'
+