ruyi 0.39.0__py3-none-any.whl

ruyi/telemetry/provider.py

@@ -0,0 +1,411 @@
+import calendar
+import datetime
+import functools
+import json
+import pathlib
+import sys
+import time
+from typing import Callable, TYPE_CHECKING, cast
+import uuid
+
+from ..log import RuyiLogger
+from .node_info import NodeInfo, gather_node_info
+from .scope import TelemetryScope
+from .store import TelemetryStore
+
+if TYPE_CHECKING:
+    # for avoiding circular import
+    from ..config import GlobalConfig
+
+FALLBACK_PM_TELEMETRY_ENDPOINT = "https://api.ruyisdk.cn/telemetry/pm/"
+
+FIRST_RUN_PROMPT = """\
+Welcome to RuyiSDK! This appears to be your first run of [yellow]ruyi[/].
+
+By default, the RuyiSDK team collects anonymous usage data to help us improve
+the product. No personal information or detail about your project is ever
+collected. The data will be uploaded to RuyiSDK team-managed servers located
+in the Chinese mainland if you agree to the uploading. You can change this
+setting at any time by running [yellow]ruyi telemetry consent[/] or
+[yellow]ruyi telemetry optout[/].
+
+We would like to ask if you agree to have basic information about this [yellow]ruyi[/]
+installation uploaded, right now, for one time. This will allow the RuyiSDK team
+to have more precise knowledge about the product's adoption. Thank you for
+your support!
+"""
+
+
+def next_utc_weekday(wday: int, now: float | None = None) -> int:
+    t = time.gmtime(now)
+    mday_delta = wday - t.tm_wday
+    if mday_delta < 0:
+        mday_delta += 7
+
+    next_t = (
+        t.tm_year,
+        t.tm_mon,
+        t.tm_mday + mday_delta,
+        0,  # tm_hour
+        0,  # tm_min
+        0,  # tm_sec
+        0,  # tm_wday
+        0,  # tm_yday
+        -1,  # tm_isdst
+    )
+    return calendar.timegm(next_t)
+
+
+def set_telemetry_mode(
+    gc: "GlobalConfig",
+    mode: str,
+    consent_time: datetime.datetime | None = None,
+    show_cli_feedback: bool = True,
+) -> None:
+    """Set telemetry mode and consent time (if applicable) in the user preference."""
+
+    from ..config.editor import ConfigEditor
+    from ..config import schema
+
+    logger = gc.logger
+
+    with ConfigEditor.work_on_user_local_config(gc) as ed:
+        ed.set_value((schema.SECTION_TELEMETRY, schema.KEY_TELEMETRY_MODE), mode)
+
+        if mode == "on":
+            if consent_time is None:
+                consent_time = datetime.datetime.now().astimezone()
+            ed.set_value(
+                (schema.SECTION_TELEMETRY, schema.KEY_TELEMETRY_UPLOAD_CONSENT),
+                consent_time,
+            )
+        else:
+            ed.unset_value(
+                (schema.SECTION_TELEMETRY, schema.KEY_TELEMETRY_UPLOAD_CONSENT)
+            )
+
+        ed.stage()
+
+    if not show_cli_feedback:
+        return
+    match mode:
+        case "on":
+            logger.I("telemetry data uploading is now enabled")
+            logger.I(
+                "you can opt out at any time by running [yellow]ruyi telemetry optout[/]"
+            )
+        case "local":
+            logger.I("telemetry mode is now set to local collection only")
+            logger.I(
+                "you can re-enable telemetry data uploading at any time by running [yellow]ruyi telemetry consent[/]"
+            )
+            logger.I(
+                "or opt out at any time by running [yellow]ruyi telemetry optout[/]"
+            )
+        case "off":
+            logger.I("telemetry data collection is now disabled")
+            logger.I(
+                "you can re-enable telemetry data uploads at any time by running [yellow]ruyi telemetry consent[/]"
+            )
+        case _:
+            raise ValueError(f"invalid telemetry mode: {mode}")
+
+
+class TelemetryProvider:
+    def __init__(self, gc: "GlobalConfig") -> None:
+        self.state_root = pathlib.Path(gc.telemetry_root)
+        self.local_mode = gc.telemetry_mode == "local"
+        self.upload_consent_time = gc.telemetry_upload_consent_time
+
+        self._discard_events = False
+        self._gc = gc
+        self._is_first_run = False
+        self._stores: dict[TelemetryScope, TelemetryStore] = {}
+        self._upload_on_exit = False
+
+        # create the PM store
+        self.init_store(TelemetryScope(None))
+        # TODO: add real multi-repo support
+        self.init_store(TelemetryScope("ruyisdk"))
+
+    @property
+    def logger(self) -> RuyiLogger:
+        return self._gc.logger
+
+    def store(self, scope: TelemetryScope) -> TelemetryStore | None:
+        return self._stores.get(scope)
+
+    def init_store(self, scope: TelemetryScope) -> None:
+        store_root = self.state_root
+        api_url_fn: Callable[[], str | None] | None = None
+        if repo_name := scope.repo_name:
+            if repo_name != "ruyisdk":
+                raise NotImplementedError("multi-repo support not implemented yet")
+            store_root = store_root / "repos" / repo_name
+
+            def _f() -> str | None:
+                # access the repo attribute lazily to speed up CLI startup
+                return self._gc.repo.get_telemetry_api_url("repo")
+
+            api_url_fn = _f
+        else:
+            # configure the PM telemetry endpoint
+            api_url_fn = functools.partial(self._detect_pm_api_url, self._gc)
+
+        store = TelemetryStore(
+            self.logger,
+            scope,
+            store_root,
+            api_url_factory=api_url_fn,
+        )
+        self._stores[scope] = store
+
+    def _detect_pm_api_url(self, gc: "GlobalConfig") -> str | None:
+        url = FALLBACK_PM_TELEMETRY_ENDPOINT
+        cfg_src = "fallback"
+        if gc.override_pm_telemetry_url is not None:
+            cfg_src = "local config"
+            url = gc.override_pm_telemetry_url
+        else:
+            if repo_provided_url := gc.repo.get_telemetry_api_url("pm"):
+                cfg_src = "repo"
+                url = repo_provided_url
+        self.logger.D(
+            f"configured PM telemetry endpoint via {cfg_src}: {url or '(n/a)'}"
+        )
+        return url
+
+    @property
+    def installation_file(self) -> pathlib.Path:
+        return self.state_root / "installation.json"
+
+    def check_first_run_status(self) -> None:
+        """Check if this is the first run of the application by checking if installation file exists.
+        This must be done before init_installation() is potentially called.
+        """
+        self._is_first_run = not self.installation_file.exists()
+
+    @property
+    def is_first_run(self) -> bool:
+        """Check if this is the first run of the application."""
+        return self._is_first_run
+
+    def init_installation(self, force_reinit: bool) -> NodeInfo | None:
+        installation_file = self.installation_file
+        if installation_file.exists() and not force_reinit:
+            return self.read_installation_data()
+
+        # either this is a fresh installation or we're forcing a refresh
+        installation_id = uuid.uuid4()
+        self.logger.D(
+            f"initializing telemetry data store, installation_id={installation_id.hex}"
+        )
+        self.state_root.mkdir(parents=True, exist_ok=True)
+
+        # (over)write installation data
+        installation_data = gather_node_info(installation_id)
+        with open(installation_file, "wb") as fp:
+            fp.write(json.dumps(installation_data).encode("utf-8"))
+        return installation_data
+
+    def read_installation_data(self) -> NodeInfo | None:
+        with open(self.installation_file, "rb") as fp:
+            return cast(NodeInfo, json.load(fp))
+
+    def upload_weekday(self) -> int | None:
+        try:
+            installation_data = self.read_installation_data()
+        except FileNotFoundError:
+            # init the node info if it's gone
+            installation_data = self.init_installation(False)
+
+        if installation_data is None:
+            return None
+
+        try:
+            report_uuid_prefix = int(installation_data["report_uuid"][:8], 16)
+        except ValueError:
+            return None
+
+        return report_uuid_prefix % 7  # 0 is Monday
+
+    def has_upload_consent(self, time_now: float | None = None) -> bool:
+        if self.upload_consent_time is None:
+            return False
+        if time_now is None:
+            time_now = time.time()
+        return self.upload_consent_time.timestamp() <= time_now
+
+    def print_telemetry_notice(self, for_cli_verbose_output: bool = False) -> None:
+        if self.local_mode:
+            if for_cli_verbose_output:
+                self.logger.I(
+                    "telemetry mode is [green]local[/]: local data collection only, no uploads"
+                )
+            return
+
+        now = time.time()
+        if self.has_upload_consent(now) and not for_cli_verbose_output:
+            self.logger.D("user has consented to telemetry upload")
+            return
+
+        upload_wday = self.upload_weekday()
+        if upload_wday is None:
+            return
+        upload_wday_name = calendar.day_name[upload_wday]
+
+        next_upload_day_ts = next_utc_weekday(upload_wday, now)
+        next_upload_day = time.localtime(next_upload_day_ts)
+        next_upload_day_end = time.localtime(next_upload_day_ts + 86400)
+        next_upload_day_str = time.strftime("%Y-%m-%d %H:%M:%S %z", next_upload_day)
+        next_upload_day_end_str = time.strftime(
+            "%Y-%m-%d %H:%M:%S %z", next_upload_day_end
+        )
+
+        today_is_upload_day = self.is_upload_day(now)
+        if for_cli_verbose_output:
+            self.logger.I(
+                "telemetry mode is [green]on[/]: data is collected and periodically uploaded"
+            )
+            self.logger.I(
+                f"non-tracking usage information will be uploaded to RuyiSDK-managed servers [bold green]every {upload_wday_name}[/]"
+            )
+        else:
+            self.logger.W(
+                f"this [yellow]ruyi[/] installation has telemetry mode set to [yellow]on[/], and [bold]will upload non-tracking usage information to RuyiSDK-managed servers[/] [bold green]every {upload_wday_name}[/]"
+            )
+
+        if today_is_upload_day:
+            for scope, store in self._stores.items():
+                has_uploaded_today = self.has_uploaded_today(scope, now)
+                if has_uploaded_today:
+                    if last_upload_time := store.last_upload_timestamp:
+                        last_upload_time_str = time.strftime(
+                            "%Y-%m-%d %H:%M:%S %z", time.localtime(last_upload_time)
+                        )
+                        self.logger.I(
+                            f"scope {scope}: usage information has already been uploaded today at {last_upload_time_str}"
+                        )
+                    else:
+                        self.logger.I(
+                            f"scope {scope}: usage information has already been uploaded sometime today"
+                        )
+                else:
+                    self.logger.I(
+                        f"scope {scope}: the next upload will happen [bold green]today[/] if not already"
+                    )
+        else:
+            self.logger.I(
+                f"the next upload will happen anytime [yellow]ruyi[/] is executed between [bold green]{next_upload_day_str}[/] and [bold green]{next_upload_day_end_str}[/]"
+            )
+
+        if not for_cli_verbose_output:
+            self.logger.I("in order to hide this banner:")
+            self.logger.I("- opt out with [yellow]ruyi telemetry optout[/]")
+            self.logger.I("- or give consent with [yellow]ruyi telemetry consent[/]")
+
+    def next_upload_day(self, time_now: float | None = None) -> int | None:
+        upload_wday = self.upload_weekday()
+        if upload_wday is None:
+            return None
+        return next_utc_weekday(upload_wday, time_now)
+
+    def is_upload_day(self, time_now: float | None = None) -> bool:
+        if time_now is None:
+            time_now = time.time()
+        if upload_day := self.next_upload_day(time_now):
+            return upload_day <= time_now
+        return False
+
+    def has_uploaded_today(
+        self,
+        scope: TelemetryScope,
+        time_now: float | None = None,
+    ) -> bool:
+        if time_now is None:
+            time_now = time.time()
+        if upload_day := self.next_upload_day(time_now):
+            upload_day_end = upload_day + 86400
+            store = self.store(scope)
+            if store is None:
+                return False
+            if last_upload_time := store.last_upload_timestamp:
+                return upload_day <= last_upload_time < upload_day_end
+        return False
+
+    def record(self, scope: TelemetryScope, kind: str, **params: object) -> None:
+        if store := self.store(scope):
+            return store.record(kind, **params)
+        self.logger.D(f"no telemetry store for scope {scope}, discarding event")
+
+    def discard_events(self, v: bool = True) -> None:
+        self._discard_events = v
+
+    def flush(self, *, upload_now: bool = False) -> None:
+        now = time.time()
+
+        # We may be self-uninstalling and purging all state data, and in this
+        # case we don't want to record anything (thus re-creating directories).
+        if self._discard_events:
+            self.logger.D("discarding collected telemetry data")
+            return
+
+        self.logger.D("flushing telemetry to persistent store")
+
+        for scope, store in self._stores.items():
+            store.persist(now)
+
+            # try to upload if forced (upload_now or _upload_on_exit), or:
+            #
+            # * we're not in local mode
+            # * today is the day
+            # * we haven't uploaded today
+            if not (upload_now or self._upload_on_exit) and (
+                self.local_mode
+                or not self.is_upload_day(now)
+                or self.has_uploaded_today(scope, now)
+            ):
+                continue
+
+            self.prepare_data_for_upload(store)
+            self.upload_staged_payloads(store)
+
+    def prepare_data_for_upload(self, store: TelemetryStore) -> None:
+        installation_data: NodeInfo | None = None
+        if store.scope.is_pm:
+            try:
+                installation_data = self.read_installation_data()
+            except FileNotFoundError:
+                # should not happen due to is_upload_day() initializing it for us
+                # beforehand, but proceed without node info nonetheless
+                pass
+
+        return store.prepare_data_for_upload(installation_data)
+
+    def upload_staged_payloads(self, store: TelemetryStore) -> None:
+        if self.local_mode:
+            return
+
+        return store.upload_staged_payloads()
+
+    def maybe_prompt_for_first_run_upload(self) -> None:
+        """
+        Ask whether the user consents to a first-run telemetry upload when
+        running for the first time (OOBE) and with an interactive stdin.
+        """
+
+        # Only prompt if this is first run and stdin is a TTY
+        if not (self.is_first_run and sys.stdin.isatty()):
+            return
+
+        from ..cli import user_input
+
+        self.logger.I(FIRST_RUN_PROMPT)
+        if not user_input.ask_for_yesno_confirmation(
+            self.logger,
+            "Do you agree?",
+            True,
+        ):
+            return
+
+        self._upload_on_exit = True

ruyi/telemetry/scope.py

@@ -0,0 +1,43 @@
+from typing import Literal, TypeAlias, TypeGuard
+
+TelemetryScopeConfig: TypeAlias = Literal["pm"] | Literal["repo"]
+
+
+def is_telemetry_scope_config(x: object) -> TypeGuard[TelemetryScopeConfig]:
+    if not isinstance(x, str):
+        return False
+    match x:
+        case "pm" | "repo":
+            return True
+        case _:
+            return False
+
+
+class TelemetryScope:
+    def __init__(self, repo_name: str | None) -> None:
+        self._repo_name = repo_name
+
+    def __repr__(self) -> str:
+        return f"TelemetryScope(repo_name={self._repo_name})"
+
+    def __str__(self) -> str:
+        if self._repo_name:
+            return f"repo:{self._repo_name}"
+        return "pm"
+
+    def __hash__(self) -> int:
+        # behave like the inner field
+        return hash(self._repo_name)
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, TelemetryScope):
+            return False
+        return self._repo_name == value._repo_name
+
+    @property
+    def repo_name(self) -> str | None:
+        return self._repo_name
+
+    @property
+    def is_pm(self) -> bool:
+        return self._repo_name is None

ruyi/telemetry/store.py

@@ -0,0 +1,238 @@
+from functools import cached_property
+import json
+import os
+import pathlib
+import re
+import time
+from typing import Callable, Final, Iterable
+import uuid
+
+from ..log import RuyiLogger
+from ..utils.url import urljoin_for_sure
+from .aggregate import UploadPayload, aggregate_events
+from .event import TelemetryEvent, is_telemetry_event
+from .node_info import NodeInfo
+from .scope import TelemetryScope
+
+# e.g. "run.202410201845.d06ca5d668e64fec833ed3e6eb926a2c.ndjson"
+RE_RAW_EVENT_FILENAME: Final = re.compile(
+    r"^run\.(?P<time_bucket>\d{12})\.(?P<uuid>[0-9a-f]{32})\.ndjson$"
+)
+
+
+def get_time_bucket(timestamp: int | float | time.struct_time | None = None) -> str:
+    if timestamp is None:
+        return time.strftime("%Y%m%d%H%M")
+    elif isinstance(timestamp, float) or isinstance(timestamp, int):
+        timestamp = time.localtime(timestamp)
+    return time.strftime("%Y%m%d%H%M", timestamp)
+
+
+def time_bucket_from_filename(filename: str) -> str | None:
+    if m := RE_RAW_EVENT_FILENAME.match(filename):
+        return m.group("time_bucket")
+    return None
+
+
+class TelemetryStore:
+    def __init__(
+        self,
+        logger: RuyiLogger,
+        scope: TelemetryScope,
+        store_root: pathlib.Path,
+        api_url: str | None = None,
+        api_url_factory: Callable[[], str | None] | None = None,
+    ) -> None:
+        self._logger = logger
+        self.scope = scope
+        self.store_root = store_root
+        self._api_url = api_url
+        self._api_url_factory = api_url_factory
+
+        self._events: list[TelemetryEvent] = []
+
+    @cached_property
+    def api_url(self) -> str | None:
+        if u := self._api_url:
+            return u
+        if f := self._api_url_factory:
+            return f()
+        return None
+
+    @property
+    def raw_events_dir(self) -> pathlib.Path:
+        return self.store_root / "raw"
+
+    @property
+    def upload_stage_dir(self) -> pathlib.Path:
+        return self.store_root / "staged"
+
+    @property
+    def uploaded_dir(self) -> pathlib.Path:
+        return self.store_root / "uploaded"
+
+    @property
+    def last_upload_marker_file(self) -> pathlib.Path:
+        return self.store_root / ".stamp-last-upload"
+
+    @property
+    def last_upload_timestamp(self) -> float | None:
+        try:
+            return self.last_upload_marker_file.stat().st_mtime
+        except FileNotFoundError:
+            return None
+
+    def record_upload_timestamp(self, time_now: float | None = None) -> None:
+        if time_now is None:
+            time_now = time.time()
+        f = self.last_upload_marker_file
+        f.touch()
+        os.utime(f, (time_now, time_now))
+
+    def record(self, kind: str, **params: object) -> None:
+        self._events.append({"fmt": 1, "kind": kind, "params": params})
+
+    def discard_events(self, v: bool = True) -> None:
+        self._discard_events = v
+
+    def persist(self, now: float | None = None) -> None:
+        if not self._events:
+            self._logger.D(f"scope {self.scope}: no event to persist")
+            return
+
+        now = time.time() if now is None else now
+
+        self._logger.D(f"scope {self.scope}: flushing telemetry to persistent store")
+
+        raw_events_dir = self.raw_events_dir
+        raw_events_dir.mkdir(parents=True, exist_ok=True)
+
+        # TODO: for now it is safe to not lock, because flush() is only ever
+        # called at program exit time
+        rough_time = get_time_bucket(now)
+        rand = uuid.uuid4().hex
+        batch_events_file = raw_events_dir / f"run.{rough_time}.{rand}.ndjson"
+        with open(batch_events_file, "wb") as fp:
+            for e in self._events:
+                payload = json.dumps(e)
+                fp.write(payload.encode("utf-8"))
+                fp.write(b"\n")
+
+        self._logger.D(
+            f"scope {self.scope}: persisted {len(self._events)} telemetry event(s)"
+        )
+
+    def upload(self, installation_data: NodeInfo | None = None) -> None:
+        self.prepare_data_for_upload(installation_data)
+        self.upload_staged_payloads()
+
+    def read_back_raw_events(self) -> Iterable[TelemetryEvent]:
+        try:
+            for f in self.raw_events_dir.glob("run.*.ndjson"):
+                time_bucket = time_bucket_from_filename(f.name)
+                with open(f, "r", encoding="utf-8", newline=None) as fp:
+                    for line in fp:
+                        try:
+                            obj = json.loads(line)
+                        except json.JSONDecodeError:
+                            # losing some malformed telemetry events is okay
+                            continue
+                        if not is_telemetry_event(obj):
+                            # ditto
+                            continue
+                        if time_bucket is not None and "time_bucket" not in obj:
+                            obj["time_bucket"] = time_bucket
+                        yield obj
+        except FileNotFoundError:
+            pass
+
+    def purge_raw_events(self) -> None:
+        files = list(self.raw_events_dir.glob("run.*.ndjson"))
+        for f in files:
+            f.unlink(missing_ok=True)
+
+    def gen_upload_staging_filename(self, nonce: str) -> pathlib.Path:
+        return self.upload_stage_dir / f"staged.{nonce}.json"
+
+    def prepare_data_for_upload(self, installation_data: NodeInfo | None) -> None:
+        # import ruyi.version here because this package is on the CLI startup
+        # critical path, and version probing is costly there
+        from ..version import RUYI_SEMVER
+
+        aggregate_data = list(aggregate_events(self.read_back_raw_events()))
+
+        payload_nonce = uuid.uuid4().hex  # for server-side dedup purposes
+        payload: UploadPayload = {
+            "fmt": 1,
+            "nonce": payload_nonce,
+            "ruyi_version": str(RUYI_SEMVER),
+            "events": aggregate_data,
+        }
+        if installation_data is not None:
+            payload["installation"] = installation_data
+
+        dest_path = self.gen_upload_staging_filename(payload_nonce)
+        self.upload_stage_dir.mkdir(parents=True, exist_ok=True)
+        dest_path.write_text(json.dumps(payload), encoding="utf-8")
+
+        self.purge_raw_events()
+
+    def upload_staged_payloads(self) -> None:
+        if not self.api_url:
+            return
+
+        try:
+            staged_payloads = list(self.upload_stage_dir.glob("staged.*.json"))
+        except FileNotFoundError:
+            return
+
+        try:
+            self.uploaded_dir.mkdir(parents=True, exist_ok=True)
+        except OSError:
+            return
+
+        for f in staged_payloads:
+            self.upload_one_staged_payload(f, self.api_url)
+
+        self.record_upload_timestamp()
+
+    def upload_one_staged_payload(
+        self,
+        f: pathlib.Path,
+        endpoint: str,
+    ) -> None:
+        # import ruyi.version here because this package is on the CLI startup
+        # critical path, and version probing is costly there
+        from ..version import RUYI_USER_AGENT
+
+        api_path = urljoin_for_sure(endpoint, "upload-v1")
+        self._logger.D(f"scope {self.scope}: about to upload payload {f} to {api_path}")
+
+        import requests
+
+        resp = requests.post(
+            api_path,
+            data=f.read_bytes(),
+            headers={"User-Agent": RUYI_USER_AGENT},
+            allow_redirects=True,
+            timeout=5,
+        )
+
+        if not (200 <= resp.status_code < 300):
+            self._logger.D(
+                f"scope {self.scope}: telemetry upload failed: status code {resp.status_code}, content {resp.content.decode('utf-8', 'replace')}"
+            )
+            return
+
+        self._logger.D(
+            f"scope {self.scope}: telemetry upload ok: status code {resp.status_code}"
+        )
+
+        # move to completed dir
+        # TODO: rotation
+        try:
+            f.rename(self.uploaded_dir / f.name)
+        except OSError as e:
+            self._logger.D(
+                f"scope {self.scope}: failed to move uploaded payload away: {e}"
+            )