ominfra 0.0.0.dev102__tar.gz → 0.0.0.dev104__tar.gz

{ominfra-0.0.0.dev102/ominfra.egg-info → ominfra-0.0.0.dev104}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ominfra
-Version: 0.0.0.dev102
+Version: 0.0.0.dev104
 Summary: ominfra
 Author: wrmsr
 License: BSD-3-Clause
@@ -12,8 +12,8 @@ Classifier: Operating System :: OS Independent
 Classifier: Operating System :: POSIX
 Requires-Python: >=3.12
 License-File: LICENSE
-Requires-Dist: omdev==0.0.0.dev102
-Requires-Dist: omlish==0.0.0.dev102
+Requires-Dist: omdev==0.0.0.dev104
+Requires-Dist: omlish==0.0.0.dev104
 Provides-Extra: all
 Requires-Dist: paramiko~=3.5; extra == "all"
 Requires-Dist: asyncssh~=2.18; extra == "all"

ominfra-0.0.0.dev104/ominfra/clouds/aws/journald2aws/cursor.py

@@ -0,0 +1,47 @@
+# ruff: noqa: UP007
+import os.path
+import typing as ta
+
+from omlish.lite.logs import log
+
+
+class JournalctlToAwsCursor:
+    def __init__(
+            self,
+            cursor_file: ta.Optional[str] = None,
+            *,
+            ensure_locked: ta.Optional[ta.Callable[[], None]] = None,
+    ) -> None:
+        super().__init__()
+        self._cursor_file = cursor_file
+        self._ensure_locked = ensure_locked
+
+    #
+
+    def get(self) -> ta.Optional[str]:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        if not (cf := self._cursor_file):
+            return None
+        cf = os.path.expanduser(cf)
+
+        try:
+            with open(cf) as f:
+                return f.read().strip()
+        except FileNotFoundError:
+            return None
+
+    def set(self, cursor: str) -> None:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        if not (cf := self._cursor_file):
+            return
+        cf = os.path.expanduser(cf)
+
+        log.info('Writing cursor file %s : %s', cf, cursor)
+        with open(ncf := cf + '.next', 'w') as f:
+            f.write(cursor)
+
+        os.rename(ncf, cf)

ominfra-0.0.0.dev104/ominfra/clouds/aws/journald2aws/driver.py

@@ -0,0 +1,210 @@
+# ruff: noqa: UP007
+"""
+TODO:
+ - create log group
+ - log stats - chunk sizes, byte count, num calls, etc
+
+==
+
+https://www.freedesktop.org/software/systemd/man/latest/journalctl.html
+
+journalctl:
+  -o json
+  --show-cursor
+
+  --since "2012-10-30 18:17:16"
+  --until "2012-10-30 18:17:16"
+
+  --after-cursor <cursor>
+
+==
+
+https://www.freedesktop.org/software/systemd/man/latest/systemd.journal-fields.html
+
+==
+
+@dc.dataclass(frozen=True)
+class Journald2AwsConfig:
+    log_group_name: str
+    log_stream_name: str
+
+    aws_batch_size: int = 1_000
+    aws_flush_interval_s: float = 1.
+"""
+import dataclasses as dc
+import os.path
+import queue
+import time
+import typing as ta
+
+from omlish.lite.cached import cached_nullary
+from omlish.lite.check import check_non_empty_str
+from omlish.lite.contextmanagers import ExitStacked
+from omlish.lite.logs import log
+from omlish.lite.pidfile import Pidfile
+from omlish.lite.runtime import is_debugger_attached
+
+from ....journald.messages import JournalctlMessage  # noqa
+from ....journald.tailer import JournalctlTailerWorker
+from ..auth import AwsSigner
+from ..logs import AwsLogMessageBuilder
+from .cursor import JournalctlToAwsCursor
+from .poster import JournalctlToAwsPosterWorker
+
+
+##
+
+
+class JournalctlToAwsDriver(ExitStacked):
+    @dc.dataclass(frozen=True)
+    class Config:
+        pid_file: ta.Optional[str] = None
+
+        cursor_file: ta.Optional[str] = None
+
+        runtime_limit: ta.Optional[float] = None
+
+        #
+
+        aws_log_group_name: str = 'omlish'
+        aws_log_stream_name: ta.Optional[str] = None
+
+        aws_access_key_id: ta.Optional[str] = None
+        aws_secret_access_key: ta.Optional[str] = dc.field(default=None, repr=False)
+
+        aws_region_name: str = 'us-west-1'
+
+        aws_dry_run: bool = False
+
+        #
+
+        journalctl_cmd: ta.Optional[ta.Sequence[str]] = None
+
+        journalctl_after_cursor: ta.Optional[str] = None
+        journalctl_since: ta.Optional[str] = None
+
+    def __init__(self, config: Config) -> None:
+        super().__init__()
+
+        self._config = config
+
+    #
+
+    @cached_nullary
+    def _pidfile(self) -> ta.Optional[Pidfile]:
+        if self._config.pid_file is None:
+            return None
+
+        pfp = os.path.expanduser(self._config.pid_file)
+
+        log.info('Opening pidfile %s', pfp)
+
+        pf = self._enter_context(Pidfile(pfp))
+        pf.write()
+        return pf
+
+    def _ensure_locked(self) -> None:
+        if (pf := self._pidfile()) is not None:
+            pf.ensure_locked()
+
+    #
+
+    @cached_nullary
+    def _cursor(self) -> JournalctlToAwsCursor:
+        return JournalctlToAwsCursor(
+            self._config.cursor_file,
+            ensure_locked=self._ensure_locked,
+        )
+
+    #
+
+    @cached_nullary
+    def _aws_credentials(self) -> ta.Optional[AwsSigner.Credentials]:
+        if self._config.aws_access_key_id is None and self._config.aws_secret_access_key is None:
+            return None
+
+        return AwsSigner.Credentials(
+            access_key_id=check_non_empty_str(self._config.aws_access_key_id),
+            secret_access_key=check_non_empty_str(self._config.aws_secret_access_key),
+        )
+
+    @cached_nullary
+    def _aws_log_message_builder(self) -> AwsLogMessageBuilder:
+        return AwsLogMessageBuilder(
+            log_group_name=self._config.aws_log_group_name,
+            log_stream_name=check_non_empty_str(self._config.aws_log_stream_name),
+            region_name=self._config.aws_region_name,
+            credentials=self._aws_credentials(),
+        )
+
+    #
+
+    @cached_nullary
+    def _journalctl_message_queue(self):  # type: () -> queue.Queue[ta.Sequence[JournalctlMessage]]
+        return queue.Queue()
+
+    @cached_nullary
+    def _journalctl_tailer_worker(self) -> JournalctlTailerWorker:
+        ac: ta.Optional[str] = None
+
+        if (since := self._config.journalctl_since):
+            log.info('Starting since %s', since)
+
+        else:
+            ac = self._config.journalctl_after_cursor
+            if ac is None:
+                ac = self._cursor().get()
+            if ac is not None:
+                log.info('Starting from cursor %s', ac)
+
+        return JournalctlTailerWorker(
+            self._journalctl_message_queue(),
+
+            since=since,
+            after_cursor=ac,
+
+            cmd=self._config.journalctl_cmd,
+            shell_wrap=is_debugger_attached(),
+        )
+
+    #
+
+    @cached_nullary
+    def _aws_poster_worker(self) -> JournalctlToAwsPosterWorker:
+        return JournalctlToAwsPosterWorker(
+            self._journalctl_message_queue(),
+            self._aws_log_message_builder(),
+            self._cursor(),
+
+            ensure_locked=self._ensure_locked,
+            dry_run=self._config.aws_dry_run,
+        )
+
+    #
+
+    def run(self) -> None:
+        pw: JournalctlToAwsPosterWorker = self._aws_poster_worker()
+        tw: JournalctlTailerWorker = self._journalctl_tailer_worker()
+
+        ws = [pw, tw]
+
+        for w in ws:
+            w.start()
+
+        start = time.time()
+
+        while True:
+            for w in ws:
+                if not w.is_alive():
+                    log.critical('Worker died: %r', w)
+                    break
+
+            if (rl := self._config.runtime_limit) is not None and time.time() - start >= rl:
+                log.warning('Runtime limit reached')
+                break
+
+            time.sleep(1.)
+
+        for w in reversed(ws):
+            w.stop()
+            w.join()

ominfra-0.0.0.dev104/ominfra/clouds/aws/journald2aws/main.py

@@ -0,0 +1,81 @@
+#!/usr/bin/env python3
+# @omlish-amalg ../../../scripts/journald2aws.py
+import argparse
+import dataclasses as dc
+import json
+import os.path
+import sys
+
+from omlish.lite.logs import configure_standard_logging
+from omlish.lite.marshal import unmarshal_obj
+
+from .driver import JournalctlToAwsDriver
+
+
+def _main() -> None:
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--config-file')
+    parser.add_argument('-v', '--verbose', action='store_true')
+
+    parser.add_argument('--after-cursor', nargs='?')
+    parser.add_argument('--since', nargs='?')
+    parser.add_argument('--dry-run', action='store_true')
+
+    parser.add_argument('--message', nargs='?')
+    parser.add_argument('--real', action='store_true')
+    parser.add_argument('--num-messages', type=int)
+    parser.add_argument('--runtime-limit', type=float)
+
+    args = parser.parse_args()
+
+    #
+
+    configure_standard_logging('DEBUG' if args.verbose else 'INFO')
+
+    #
+
+    config: JournalctlToAwsDriver.Config
+    if args.config_file:
+        with open(os.path.expanduser(args.config_file)) as cf:
+            config_dct = json.load(cf)
+        config = unmarshal_obj(config_dct, JournalctlToAwsDriver.Config)
+    else:
+        config = JournalctlToAwsDriver.Config()
+
+    #
+
+    for k in ['aws_access_key_id', 'aws_secret_access_key']:
+        if not getattr(config, k) and k.upper() in os.environ:
+            config = dc.replace(config, **{k: os.environ.get(k.upper())})  # type: ignore
+
+    #
+
+    if not args.real:
+        config = dc.replace(config, journalctl_cmd=[
+            sys.executable,
+            os.path.join(os.path.dirname(__file__), '..', '..', '..', 'journald', 'genmessages.py'),
+            '--sleep-n', '2',
+            '--sleep-s', '.5',
+            *(['--message', args.message] if args.message else []),
+            str(args.num_messages or 100_000),
+        ])
+
+    #
+
+    for ca, pa in [
+        ('journalctl_after_cursor', 'after_cursor'),
+        ('journalctl_since', 'since'),
+        ('aws_dry_run', 'dry_run'),
+    ]:
+        if (av := getattr(args, pa)):
+            config = dc.replace(config, **{ca: av})
+
+    #
+
+    with JournalctlToAwsDriver(config) as jta:
+        jta.run()
+
+
+if __name__ == '__main__':
+    _main()

ominfra-0.0.0.dev104/ominfra/clouds/aws/journald2aws/poster.py

@@ -0,0 +1,89 @@
+# ruff: noqa: UP007
+"""
+TODO:
+ - retries
+"""
+import json
+import queue
+import time
+import typing as ta
+import urllib.request
+
+from omlish.lite.logs import log
+
+from ....journald.messages import JournalctlMessage  # noqa
+from ....threadworkers import ThreadWorker
+from ..logs import AwsLogMessageBuilder
+from ..logs import AwsPutLogEventsResponse
+from .cursor import JournalctlToAwsCursor
+
+
+class JournalctlToAwsPosterWorker(ThreadWorker):
+    def __init__(
+            self,
+            queue,  # type: queue.Queue[ta.Sequence[JournalctlMessage]]  # noqa
+            builder: AwsLogMessageBuilder,
+            cursor: JournalctlToAwsCursor,
+            *,
+            ensure_locked: ta.Optional[ta.Callable[[], None]] = None,
+            dry_run: bool = False,
+            queue_timeout_s: float = 1.,
+            **kwargs: ta.Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._queue = queue
+        self._builder = builder
+        self._cursor = cursor
+        self._ensure_locked = ensure_locked
+        self._dry_run = dry_run
+        self._queue_timeout_s = queue_timeout_s
+    #
+
+    def _run(self) -> None:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        last_cursor: ta.Optional[str] = None  # noqa
+        while True:
+            self._heartbeat()
+
+            try:
+                msgs: ta.Sequence[JournalctlMessage] = self._queue.get(timeout=self._queue_timeout_s)
+            except queue.Empty:
+                msgs = []
+
+            if not msgs:
+                log.debug('Empty queue chunk')
+                continue
+
+            log.debug('%r', msgs)
+
+            cur_cursor: ta.Optional[str] = None
+            for m in reversed(msgs):
+                if m.cursor is not None:
+                    cur_cursor = m.cursor
+                    break
+
+            feed_msgs = []
+            for m in msgs:
+                feed_msgs.append(AwsLogMessageBuilder.Message(
+                    message=json.dumps(m.dct, sort_keys=True),
+                    ts_ms=int((m.ts_us / 1000.) if m.ts_us is not None else (time.time() * 1000.)),
+                ))
+
+            for post in self._builder.feed(feed_msgs):
+                log.debug('%r', post)
+
+                if not self._dry_run:
+                    with urllib.request.urlopen(urllib.request.Request(  # noqa
+                            post.url,
+                            method='POST',
+                            headers=dict(post.headers),
+                            data=post.data,
+                    )) as resp:
+                        response = AwsPutLogEventsResponse.from_aws(json.loads(resp.read().decode('utf-8')))
+                    log.debug('%r', response)
+
+            if cur_cursor is not None:
+                self._cursor.set(cur_cursor)
+                last_cursor = cur_cursor  # noqa

{ominfra-0.0.0.dev102 → ominfra-0.0.0.dev104}/ominfra/clouds/aws/logs.py

@@ -62,7 +62,7 @@ class AwsPutLogEventsResponse(AwsDataclass):
 ##
 
 
-class AwsLogMessagePoster:
+class AwsLogMessageBuilder:
     """
     TODO:
      - max_items
@@ -88,7 +88,7 @@ class AwsLogMessagePoster:
             log_group_name: str,
             log_stream_name: str,
             region_name: str,
-            credentials: AwsSigner.Credentials,
+            credentials: ta.Optional[AwsSigner.Credentials],
 
             url: ta.Optional[str] = None,
             service_name: str = DEFAULT_SERVICE_NAME,
@@ -110,11 +110,16 @@ class AwsLogMessagePoster:
             headers = {**headers, **extra_headers}
         self._headers = {k: [v] for k, v in headers.items()}
 
-        self._signer = V4AwsSigner(
-            credentials,
-            region_name,
-            service_name,
-        )
+        signer: ta.Optional[V4AwsSigner]
+        if credentials is not None:
+            signer = V4AwsSigner(
+                credentials,
+                region_name,
+                service_name,
+            )
+        else:
+            signer = None
+        self._signer = signer
 
     #
 
@@ -158,13 +163,14 @@ class AwsLogMessagePoster:
             payload=body,
         )
 
-        sig_headers = self._signer.sign(
-            sig_req,
-            sign_payload=False,
-        )
-        sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
+        if (signer := self._signer) is not None:
+            sig_headers = signer.sign(
+                sig_req,
+                sign_payload=False,
+            )
+            sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
 
-        post = AwsLogMessagePoster.Post(
+        post = AwsLogMessageBuilder.Post(
             url=self._url,
             headers={k: check_single(v) for k, v in sig_req.headers.items()},
             data=sig_req.payload,

ominfra-0.0.0.dev104/ominfra/clouds/gcp/auth.py

@@ -0,0 +1,48 @@
+import json
+import time
+import typing as ta
+
+from omlish import check
+from omlish import http
+from omlish.http import jwt
+
+
+DEFAULT_JWT_SCOPE = 'https://www.googleapis.com/auth/cloud-platform'
+
+
+def generate_gcp_jwt(
+        creds_dct: ta.Mapping[str, ta.Any],
+        *,
+        issued_at: int | None = None,
+        lifetime_s: int = 3600,
+        scope: str = DEFAULT_JWT_SCOPE,
+) -> str:
+    return jwt.generate_jwt(
+        issuer=creds_dct['client_email'],
+        subject=creds_dct['client_email'],
+        audience=creds_dct['token_uri'],
+        issued_at=(issued_at := int(issued_at if issued_at is not None else time.time())),
+        expires_at=issued_at + lifetime_s,
+        scope=scope,
+        key=creds_dct['private_key'],
+        algorithm='RS256',
+    )
+
+
+def get_gcp_access_token(
+        creds_dct: ta.Mapping[str, ta.Any],
+        *,
+        client: http.HttpClient | None = None,
+) -> str:
+    signed_jwt = generate_gcp_jwt(creds_dct)
+    resp = http.request(
+        creds_dct['token_uri'],
+        'POST',
+        data=jwt.build_get_token_body(signed_jwt).encode('utf-8'),
+        headers={
+            http.consts.HEADER_CONTENT_TYPE: http.consts.CONTENT_TYPE_FORM_URLENCODED,
+        },
+        client=client,
+    )
+    resp_dct = json.loads(check.not_none(resp.data).decode('utf-8'))
+    return resp_dct['access_token']