Edit

[EDIT] FILE: license.py

import asyncio
import base64
import binascii
import datetime
import json
import os
import shutil
import subprocess
import tempfile
import time
from contextlib import suppress
from json import JSONDecodeError
from pathlib import Path
from subprocess import TimeoutExpired
from typing import Optional

from peewee import OperationalError

from defence360agent.application.determine_hosting_panel import (
    is_cpanel_installed,
)
from defence360agent.contracts import sentry
from defence360agent.contracts.config import (
    ANTIVIRUS_MODE,
    Core,
    CustomBilling,
    int_from_envvar,
    logger,
)
from defence360agent.contracts.hook_events import HookEvent
from defence360agent.internals.global_scope import g
from defence360agent.subsys.panels.plesk.upgrade_urls import (
    get_plesk_upgrade_urls,
)
from defence360agent.utils import retry_on, timed_cache
from defence360agent.utils.common import HOUR, rate_limit
from defence360agent.utils.ipecho import APIError, IPEchoAPI
from defence360agent.utils.validate import IP

AV_DEFAULT_ID = "IMUNIFYAV"
UNLIMITED_USERS_COUNT = 2147483647

# no need to check the license file more often than
# once every 10 minutes, this should be enough to fix DEF-14677
_CACHE_LICENSE_TOKEN_TIMEOUT = int_from_envvar(
    "IMUNIFY360_CACHE_LICENSE_TOKEN_TIMEOUT",
    10 * 60,  # in seconds
)
# path to openssl binary used to check license signature
# we need to check several paths because of different OSes
# and different installation paths with fallback to system default
if not (OPENSSL_BIN := Path("/opt/alt/openssl11/bin/openssl")).exists():
    if not (OPENSSL_BIN := Path("/opt/alt/openssl/bin/openssl")).exists():
        OPENSSL_BIN = Path("/usr/bin/openssl")

throttled_log_error = rate_limit(period=HOUR, on_drop=logger.warning)(
    logger.error
)
throttled_log_no_v2 = rate_limit(period=HOUR, on_drop=logger.warning)(
    logger.error
)

class LicenseError(Exception):
    """Used to communicate that some function requires a license"""

class LicenseCLN:
    VERIFY_FIELDS_V1 = (
        "id",
        "status",
        "group",
        "limit",
        "token_created_utc",
        "token_expire_utc",
    )

VERIFY_FIELDS_V2 = (
        "id",
        "status",
        "limit",
        "token_created_utc",
        "token_expire_utc",
        "group_id",
        "permissions",
    )

VERIFY_FIELDS_MAP = {
        1: VERIFY_FIELDS_V1,
        2: VERIFY_FIELDS_V2,
    }

_PUBKEY_FILE = "/usr/share/imunify360/cln-pub.key"
    _ALTERNATIVE_PUBKEY_FILES = (
        # keys for self-signed licenses
        "/usr/share/imunify360/alt-license-pub.key",
    )
    _LICENSE_FILE = "/var/imunify360/license.json"
    _FREE_LICENSE_FILE = "/var/imunify360/license-free.json"
    AV_PLUS_BUY_URL = (
        "https://cln.cloudlinux.com/console/purchase/ImunifyAvPlus"
    )
    IM360_BUY_URL_TEMPLATE = (
        "https://www.cloudlinux.com/upgrade-imunify-{user_count}/"
    )
    CPANEL_UPGRADE_URL = (
        "../../../scripts14/purchase_imunifyavplus_init_IMUNIFY"
    )
    CPANEL_UPGRADE_URL_360 = (
        "../../../scripts14/purchase_imunify360_init_IMUNIFY"
    )

VERSION_THRESHOLDS = [1, 30, 250]

_token = {}
    users_count = None

@staticmethod
    @retry_on(TimeoutExpired, max_tries=2)
    def _verify_signature(
        pubkey_path: str, content: bytes, signature: bytes
    ) -> tuple[bool, Optional[list[str]]]:
        """Verify that `content` is correctly signed with public key from file
        `pubkey_path` with resulting `signature`. Returns a tuple with (success, error_list).
        """
        errors: list[str] = []
        result = False

with tempfile.NamedTemporaryFile(delete=True) as sig_file:
            sig_file.write(signature)
            sig_file.flush()
            cmd = [
                OPENSSL_BIN,
                "dgst",
                "-sha512",
                "-verify",
                pubkey_path,
                "-signature",
                sig_file.name,
            ]
            try:
                p = subprocess.run(
                    cmd,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    input=content,
                    timeout=5,
                )
            except FileNotFoundError as e:
                errors.append(f"openssl command failed: missing {e.filename}")
            else:
                if p.returncode == 0:
                    result = True
                else:
                    errors.append(
                        "Signature verification failed - "
                        f"openssl returned {p.returncode}. "
                        f"stdout: {p.stdout}, stderr: {p.stderr}"
                    )

return result, errors or None

@classmethod
    def _get_signature_input(cls, license, version: int = 1) -> bytes:
        parts = []
        for key in cls.VERIFY_FIELDS_MAP[version]:
            value = license[key]
            if isinstance(value, dict):
                parts.append(
                    "".join(
                        f"{subkey}={subvalue}"
                        for subkey, subvalue in value.items()
                    )
                )
            elif value is None:
                parts.append("null")
            else:
                parts.append(str(value))
        return "".join(parts).encode()

@classmethod
    def _find_signature(
        cls, license_token, signature_list: list[tuple[str, int]]
    ) -> tuple[Optional[str], bool]:
        """
        Verify signatures in license

:return: signature, is_alternative, version
        """
        sign: str
        all_errors: list[str] = []

def verify_and_collect_errors(*args, **kwargs):
            success, errors = cls._verify_signature(*args, **kwargs)
            if errors:
                all_errors.extend(errors)
            return success

for sign, version in signature_list:
            signature = base64.b64decode(sign)

try:
                content = cls._get_signature_input(
                    license_token, version=version
                )
            except KeyError:
                continue

if verify_and_collect_errors(cls._PUBKEY_FILE, content, signature):
                return sign, False

for alt_pubkey in cls._ALTERNATIVE_PUBKEY_FILES:
                if verify_and_collect_errors(alt_pubkey, content, signature):
                    return sign, True

for error in all_errors:
            logger.warning("%s", error)

return None, False

@classmethod
    def _load_token(cls, path):
        """
        Load license token from file and verify signature
        If signature verification successful, put
        first valid signature to 'sign' field of license
        token

:return: license token
        """
        default = {}  # default value returned on error
        try:
            with open(path) as f:
                license_token = json.load(f)

if not isinstance(license_token, dict):
                    logger.error(
                        "Failed to load license. Expected JSON object, got %r"
                        % (license_token,)
                    )
                    return default

signature, is_alternative = cls._find_signature(
                    license_token,
                    [
                        (sign, 1)
                        for sign in license_token.get("signatures", [])
                    ],
                )
                v2_sign = license_token.get("signature_v2")
                if v2_sign:
                    _sign, _ = cls._find_signature(
                        license_token, [(v2_sign, 2)]
                    )
                    if _sign is None:
                        throttled_log_error(
                            "Failed to verify license signature v2"
                        )
                        license_token.pop("permissions", None)
                elif "permissions" in license_token:
                    license_token.pop("permissions")
                    throttled_log_no_v2(
                        "License missing signature_v2 but contained "
                        "permissions; stripped (possible tampering or "
                        "stale token)"
                    )

if signature is None:
                    throttled_log_error("Failed to verify license signature")
                    return default

license_token["sign"] = signature
                license_token["is_alternative"] = is_alternative
                return license_token

except FileNotFoundError:
            # this is a common case
            logger.info("Failed to load license: not registered?")
        except JSONDecodeError as e:
            # Likely a TOCTOU read of the file mid-write by the updater;
            # the next read cycle will pick up the fully-written content.
            logger.warning("Failed to load license: %s", e)
        except (
            OSError,
            KeyError,
            UnicodeDecodeError,
            binascii.Error,
            TypeError,
        ) as e:
            # not loading broken license
            logger.error("Failed to load license: %s", e)
        return default

@classmethod
    @timed_cache(
        datetime.timedelta(seconds=_CACHE_LICENSE_TOKEN_TIMEOUT), maxsize=1
    )
    def get_token(cls) -> dict:
        """
        Get available license.
        In Antivirus mode, if main license is unavailable, return free license

:return: license token
        """
        lic_token = {}
        license_files = (
            [cls._LICENSE_FILE, cls._FREE_LICENSE_FILE]
            if ANTIVIRUS_MODE
            else [cls._LICENSE_FILE]
        )
        for lf in license_files:
            lic_token = cls._load_token(lf)
            if lic_token:
                return lic_token
        return lic_token

@classmethod
    def get_server_id(cls) -> Optional[str]:
        """
        :return: server id
        """
        return cls.get_token().get("id")

@classmethod
    def is_registered(cls):
        """
        :return: bool: if we have token
        """
        return bool(cls.get_token())

@classmethod
    def is_valid_av_plus(cls):
        """
        :return: Return true only if we have valid ImunifyAV+ or
        Imunify360 license
        """
        return ANTIVIRUS_MODE and cls.is_valid() and (not cls.is_free())

@classmethod
    def is_free(cls):
        if not ANTIVIRUS_MODE:
            return False
        return cls.get_server_id() == AV_DEFAULT_ID

@classmethod
    def is_cloud_assisted_cleanup_allowed(cls) -> bool:
        """Cloud-assisted cleanup is a paid feature (Imunify360 / ImunifyAV+)."""
        if not ANTIVIRUS_MODE:
            return True
        return cls.is_valid_av_plus()

@classmethod
    def is_valid(cls, token=None):
        """License check based on license token

return True - if license token is valid for this server
        return False - if license token is invalid
        """
        token = token or cls.get_token()
        if not token:
            return False

if ANTIVIRUS_MODE:
            return (
                token.get("status", "").startswith("ok")
                and token["token_expire_utc"] >= time.time()
            )

return (
            token["status"] in ("ok", "ok-trial")
            and token["token_expire_utc"] >= time.time()
            and (cls.users_count is None or cls.users_count <= token["limit"])
        )

@classmethod
    def has_permission(cls, permission: str, token=None):
        """License check for a specific permission based on a license token

return True - if license token has a given permission for this server
        return False - if license token does not have permission
        """
        token = token or cls.get_token()
        if not token:
            return False

return (
            permission in (perm := token.get("permissions", {}))
            and perm[permission] == "ENABLED"
        )

@classmethod
    def update(cls, token):
        """
        Write new license token to file
        :param token: new token
        :return:
        """

old_token = cls.get_token()

# Save user_limit under different name only during registration
        # Check if this is initial registration by checking if old token exists
        if not old_token and token.get("limit") is not None:
            token["saved_user_limit"] = token["limit"]

temp_file = cls._LICENSE_FILE + ".tmp"
        flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL
        mode = 0o640

with suppress(FileNotFoundError):
            os.unlink(temp_file)
        with os.fdopen(os.open(temp_file, flags, mode), "w") as f:
            json.dump(token, f)

shutil.chown(temp_file, user="root", group="_imunify")
        os.rename(temp_file, cls._LICENSE_FILE)
        cls.get_token.cache_clear()
        sentry.set_server_id(cls.get_server_id())
        sentry.set_product_name(cls.get_product_name())
        try:
            cls.renew_hook(old_token, token)
        except OperationalError:
            pass

@classmethod
    def renew_hook(cls, old_token, token):
        important_keys = ["license_expire_utc", "status", "limit", "id"]
        exp_time = token.get("license_expire_utc")
        license_type = cls.fill_license_type(token)
        condition = any(
            [token.get(elem) != old_token.get(elem) for elem in important_keys]
        )

if condition:
            license_updated = HookEvent.LicenseRenewed(
                exp_time=exp_time, license=license_type
            )
            from defence360agent.hooks.execute import execute_hooks

asyncio.gather(
                execute_hooks(license_updated), return_exceptions=True
            )

@classmethod
    def delete(cls):
        """
        Delete license token along with old-style license data
        :return:
        """
        with suppress(FileNotFoundError):
            os.unlink(cls._LICENSE_FILE)
        cls.get_token.cache_clear()
        sentry.set_server_id(None)
        sentry.set_product_name(cls.get_product_name())

@classmethod
    def fill_license_type(cls, token):
        license_type = token.get("status")
        license_type_to_product = {
            "ok": "imunify360",
            "ok-trial": "imunify360Trial",
            "ok-av": "imunifyAV",
            "ok-avp": "imunifyAVPlus",
        }
        return license_type_to_product.get(license_type)

@classmethod
    def get_license_type(cls):
        return cls.fill_license_type(cls.get_token())

@classmethod
    def is_ip_license_type(cls):
        token = cls.get_token()
        if token.get("id", "").lower().startswith("ip-"):
            return True
        return False

@classmethod
    def format_upgrade_url(cls, url_template: Optional[str]) -> Optional[str]:
        """Format upgrade URL template with available parameters.

Args:
            url_template: URL template string that may contain
                {user_count}, {iaid}, and {users} placeholders

Returns:
            Formatted URL with placeholders replaced with actual values
        """
        if not url_template:
            return url_template

n = cls.users_count
        iaid = g.get("iaid", "")

# Determine user_count value
        if n is None:
            user_count = 1
        else:
            for threshold in cls.VERSION_THRESHOLDS:
                if n <= threshold:
                    user_count = threshold
                    break
            else:
                user_count = "unlimited"

url_template = url_template.replace("{user_count}", str(user_count))
        url_template = url_template.replace("{iaid}", iaid)
        url_template = url_template.replace(
            "{users}", str(n if n is not None else 1)
        )
        return url_template

@classmethod
    def _get_license_tier_recommendation(cls, user_count):
        """Get recommended license tier based on user count."""
        if user_count == 1:
            return "Single user"
        elif user_count <= 30:
            return "Up to 30 users"
        elif user_count <= 250:
            return "Up to 250 users"
        else:
            return "Unlimited users"

@classmethod
    def _format_license_exceeded_message(cls, user_count):
        """Format enhanced message when user count exceeds saved limit."""
        if user_count is None:
            return (
                "WARNING: License is invalid for current server. "
                "Unable to determine user count; please check KB article: "
                "https://cloudlinux.zendesk.com/hc/en-us/articles/"
            )

tier_name = cls._get_license_tier_recommendation(user_count)
        user_word = "user" if user_count == 1 else "users"

return (
            "WARNING: License is invalid for current server. "
            f"Detected {user_count} {user_word} → "
            f'purchase the "{tier_name}" Imunify360 license. '
            "Pricing: https://imunify360.com/pricing"
        )

@classmethod
    def license_info(cls):
        token = cls.get_token()
        key_360 = token.get("status") in ("ok", "ok-trial")

message = token.get("message", None)
        if (
            ANTIVIRUS_MODE
            and CustomBilling.UPGRADE_URL
            and not CustomBilling.NOTIFICATIONS
        ):
            message = None
        if ANTIVIRUS_MODE and key_360 and not message:
            # TODO: remove after auto-upgrade will be implemented
            message = (
                "You've got a license for the advanced security product "
                "Imunify360. Please, uninstall ImunifyAV and replace it with "
                "the Imunify360 providing comprehensive security for your "
                "server. Here are the steps for upgrade: "
                "https://docs.imunify360.com/installation/"
            )

if token:
            info = {
                "status": cls.is_valid(),
                "expiration": token.get("license_expire_utc", 0),
                "user_limit": token.get("limit"),
                "id": token.get("id"),
                "user_count": cls.users_count,
                "message": message,
                "license_type": cls.fill_license_type(token),
            }

# Generate enhanced message if license is invalid due to user limit exceeded
            # Compare against saved_user_limit instead of current limit
            if (
                not ANTIVIRUS_MODE
                and token.get("saved_user_limit") is not None
                and cls.users_count is not None
                and cls.users_count > token.get("saved_user_limit")
            ):
                info["message"] = cls._format_license_exceeded_message(
                    cls.users_count
                )
        else:
            info = {"status": False}

info["upgrade_url"] = None
        info["upgrade_url_360"] = None
        if ANTIVIRUS_MODE:
            ignored_messages = [
                "user limits",
            ]
            if info.get("message"):
                for msg in ignored_messages:
                    if msg in info["message"]:
                        info["message"] = None

# Only add ip_license when we have a valid token, since the schema
            # for status=false doesn't allow this property (additionalProperties: false)
            if token:
                info["ip_license"] = CustomBilling.IP_LICENSE and (
                    CustomBilling.UPGRADE_URL is not None
                    or CustomBilling.UPGRADE_URL_360 is not None
                )
            plesk_urls = get_plesk_upgrade_urls()
            info["upgrade_url"] = (
                cls.format_upgrade_url(CustomBilling.UPGRADE_URL)
                or plesk_urls["buy_url"]
                or token.get("upgrade_url")
                or cls.AV_PLUS_BUY_URL
            )
            info["upgrade_url_360"] = (
                cls.format_upgrade_url(CustomBilling.UPGRADE_URL_360)
                or plesk_urls["upgrade_license_url"]
                or plesk_urls["buy_url"]
                or upgrade_url_default()
            )
        # redirect_url is required for the no-license schema (status=false)
        # Set it to None when there's no token, and from token otherwise
        if not token:
            info["redirect_url"] = None
        elif not ANTIVIRUS_MODE:
            info["redirect_url"] = token.get("upgrade_url", None)
        if cls.is_demo():  # pragma: no cover
            info["demo"] = True

info[
            "eligible_for_imunify_patch"
        ] = cls.is_eligible_for_imunify_patch()

return info

@classmethod
    def is_vps(cls) -> bool:
        return cls.users_count is not None and cls.users_count <= 1

@classmethod
    def is_custom_reseller_configured(cls) -> bool:
        upgrade_urls: list[Optional[str]] = [None]
        upgrade_urls_360: list[Optional[str]] = [None]

if is_cpanel_installed():
            upgrade_urls.append(cls.CPANEL_UPGRADE_URL)
            upgrade_urls_360.append(cls.CPANEL_UPGRADE_URL_360)

# customer has any upgrade url other than default ones
        return not (
            CustomBilling.UPGRADE_URL in upgrade_urls
            and CustomBilling.UPGRADE_URL_360 in upgrade_urls_360
        )

@classmethod
    def is_eligible_for_imunify_patch(cls) -> bool:
        return (
            cls.is_vps()
            and cls.is_free()
            and not cls.is_custom_reseller_configured()
        )

@classmethod
    def get_product_name(cls) -> str:
        if not ANTIVIRUS_MODE:
            return Core.NAME

license_status = cls.get_token().get("status", "")

if license_status == "ok-av":
            return "imunify.av"
        elif license_status in ("ok-avp", "ok", "ok-trial"):
            return "imunify.av+"
        else:
            logger.error("Unknown license %s", license_status)
            return "Unknown license"

@classmethod
    def is_demo(cls) -> bool:
        return os.path.isfile("/var/imunify360/demo")

@classmethod
    def is_unlimited(cls):
        token = cls.get_token()
        return token.get("limit", 0) >= UNLIMITED_USERS_COUNT

@classmethod
    def get_im360_buy_url(cls) -> str:
        if cls.users_count is None:
            return cls.IM360_BUY_URL_TEMPLATE.format(user_count=1)
        for threshold in cls.VERSION_THRESHOLDS:
            if cls.users_count <= threshold:
                return cls.IM360_BUY_URL_TEMPLATE.format(user_count=threshold)
        return cls.IM360_BUY_URL_TEMPLATE.format(user_count="unlimited")

def upgrade_url_default():
    n = LicenseCLN.users_count
    iaid = g.get("iaid", "")

if (
        # apply custom direct store links on cPanel
        is_cpanel_installed()
        # where upgrade URL is not set or set to the old value
        and CustomBilling.UPGRADE_URL == LicenseCLN.CPANEL_UPGRADE_URL
    ):
        # We have a complex default value for cPanel installations configured
        # with that use cPanel as a reseller. They don't populate
        # the CUSTOM_BILLING config well, so we generate the links here.
        # (they care less about upsell than we do)

if not _eligible_for_new_upgrade_links(iaid):
            # A/B experiment control (old) variant:
            # return the old URL that leads through cPanel login page
            return LicenseCLN.CPANEL_UPGRADE_URL_360

# A/B experiment test (new) variant:
        # return the new URL that leads directly to the store
        base_url = (
            "https://store.cpanel.net/index.php?rp=/store/partner-addons/"
        )
        server_ip = ""
        try:
            ip = IPEchoAPI.server_ip()
            # cPanel Store only accepts w4, not IPv6
            # If we got IPv6, leave the IP field blank
            if IP.is_valid_ipv4_addr(ip):
                server_ip = ip
            else:
                logger.info(
                    "Server IP is IPv6 (%s), cPanel Store requires IPv4. "
                    "Omitting IP parameter.",
                    ip,
                )
        except APIError as e:
            logger.warning("Failed to get server IP: %s", e)

if n == 1:
            suffix = (
                f"imunify360-for-cpanel-solo&customfield%5B55%5D={server_ip}"
            )
        else:
            suffix = f"imunify360&customfield%5B375%5D={server_ip}"
        return base_url + suffix

return (
        LicenseCLN.get_im360_buy_url()
        + f"?iaid={iaid}"
        + f"&users={n}" * bool(n)
    )

def _eligible_for_new_upgrade_links(iaid: str) -> bool:
    logger.debug("checking if iaid: %s is eligible for upgrade", iaid)
    if len(iaid) == 0:
        logger.warning("receive empty iaid, fallback to old link")
        return False
    try:
        hex_bucket = int(iaid[0], 16)
    except ValueError:
        logger.warning("iaid is not hex, fallback to old link")
        return False
    hex_mid = 8

# check if iaid falls under 50% of iaid distribution
    return hex_bucket < hex_mid

CANCEL