{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "tee.ipynb"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# trusted execution environment (TEE) \u2014 Python demo\n\nNumerical companion to the entry [trusted execution environment (TEE)](https://dictionaryofml.org/terms/tee.html) of the [Dictionary of Applied Machine Learning](https://dictionaryofml.org/): it recomputes what the entry states and prints one line per check.\n\nModels the entry's Raspberry Pi in a clinic: a secure world that holds the model parameters of a learned hypothesis and the recorded heartbeats, a normal world that sees only the predictions, remote attestation by a signed digest of the code, and the two things a TEE does not protect by itself -- what the predictions reveal, and the side channel of a power trace. Self-contained (numpy/matplotlib and the standard library only), fixed seed.\n\nRequires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Generated from [`pythondemos/tee.py`](https://dictionaryofml.org/terms/tee.py); CC BY 4.0."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Notebook shim: the script resolves output paths relative to __file__,\n# which a notebook kernel does not define; everything lands in the\n# working directory instead.\nimport os\n__file__ = os.path.join(os.getcwd(), \"tee.py\")\nos.makedirs(\"pythondemos\", exist_ok=True)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "\"\"\"\ntee.py \u2014 numerical companion to the glossary entry 'trusted execution\nenvironment (TEE)'.\n\nPurpose\n-------\nModels the entry's Raspberry Pi in a clinic: a secure world that holds\nthe model parameters of a learned hypothesis and the recorded heartbeats,\na normal world that sees only the predictions, remote attestation by a\nsigned digest of the code, and the two things a TEE does not protect by\nitself -- what the predictions reveal, and the side channel of a power\ntrace.  Self-contained (numpy/matplotlib and the standard library only),\nfixed seed.\n\nSetup\n-----\nHeartbeats with eight features (interval to the previous beat, width\nof the QRS complex, height of the R peak, height of the T wave, and four\nfurther measurements that carry no information about the label) and the\nlabel arrhythmia / normal, generated by a fixed rule with noise.  The\nlearned hypothesis assigns a score w^T x to a heartbeat; its model\nparameters were obtained by training on 20 recorded heartbeats, and the\nprediction is the sign of the score.\n\nBlocks\n------\n[B-isolation]   The secure world answers prediction requests; the normal\n                world's interface hands back the prediction and nothing\n                else, so the model parameters and the heartbeats are\n                not among the values it can obtain.\n[B-attestation] The digest of the prediction code, signed with the\n                vendor key, verifies on the client; an altered code that\n                flags fewer arrhythmias changes the digest and fails\n                attestation.\n[B-reveal]      What the prediction reveals is not protected: on the 20\n                recorded heartbeats every score has the sign of the label\n                and a large size, while on new heartbeats the scores are\n                smaller and sometimes wrong, so a membership inference\n                attack that only observes predictions and knows the label\n                tells recorded from new heartbeats well above chance.\n[B-sidechannel] A power trace modeled as the number of one-bits of each\n                operand in the multiplications of w^T x leaks the model\n                parameters: the traces of different heartbeats differ with\n                the feature values, and from them the sign and size of\n                each weight is recovered; constant-time code processes\n                every operand the same way and gives a flat trace.\n\nOutputs\n-------\ntee_heartbeats.csv : the features, label and prediction of the heartbeats\n                     the normal world can see (predictions only)\ntee_scores.csv     : signed scores y w^T x of recorded and of new heartbeats\ntee_trace.csv      : the power trace of the data-dependent and of the\n                     constant-time code for one heartbeat\ntee.png            : matplotlib preview (checking only)\n\"\"\"\n\nimport hashlib\nimport hmac\nimport inspect\n\nimport numpy as np\nimport matplotlib\n\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\n\nfrom pathlib import Path\n\nOUT_DIR = Path(__file__).parent\n\nreport = []\n\n\ndef check(name, ok):\n    report.append((name, bool(ok)))\n    print(f\"  [{'ok' if ok else 'FAIL'}] {name}\")\n\n\nrng = np.random.default_rng(7)\nN_REC, N_NEW = 20, 200\nFEATURES = [\"beat interval (s)\", \"QRS width (ms)\", \"R height (mV)\", \"T height (mV)\",\n            \"P height (mV)\", \"ST level (mV)\", \"RR variability (s)\", \"noise level (mV)\"]\n\n\ndef heartbeats(n):\n    \"\"\"Features and labels (1 = arrhythmia, -1 = normal) from a fixed model.\"\"\"\n    y = np.where(rng.random(n) < 0.3, 1, -1)\n    x = np.c_[0.85 + 0.12 * rng.standard_normal(n) + 0.25 * (y == 1) * rng.standard_normal(n),\n              95 + 12 * rng.standard_normal(n) + 30 * (y == 1),\n              1.1 + 0.2 * rng.standard_normal(n) - 0.3 * (y == 1),\n              0.3 + 0.1 * rng.standard_normal(n) + 0.15 * (y == 1),\n              0.15 + 0.05 * rng.standard_normal(n), 0.05 * rng.standard_normal(n),\n              0.04 + 0.02 * rng.standard_normal(n), 0.02 + 0.01 * rng.standard_normal(n)]\n    return x, y\n\n\nx_rec, y_rec = heartbeats(N_REC)\nx_new, y_new = heartbeats(N_NEW)\nmean, sd = x_rec.mean(0), x_rec.std(0)\n\n\ndef standardize(x):\n    return (x - mean) / sd\n\n\n# the learned hypothesis: a score w^T x with model parameters trained on the recorded heartbeats\ndef train(x, y, steps=20000, lrate=0.5):\n    w = np.zeros(x.shape[1] + 1); xb = np.c_[standardize(x), np.ones(len(x))]\n    for _ in range(steps):\n        p = 1 / (1 + np.exp(-y * (xb @ w)))\n        w += lrate * (xb.T @ ((1 - p) * y)) / len(y)\n    return w\n\n\nw_hat = train(x_rec, y_rec)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-isolation]** The secure world answers prediction requests; the normal world's interface hands back the prediction and nothing else, so the model parameters and the heartbeats are not among the values it can obtain."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "class SecureWorld:\n    \"\"\"Holds the model parameters and the heartbeats; exposes predictions only.\"\"\"\n\n    def __init__(self, w, code):\n        self.__w = w.copy(); self.__beats = []; self.__code = code\n\n    def predict(self, x):\n        self.__beats.append(x)\n        return self.__code(self.__w, standardize(np.atleast_2d(x))[0])\n\n    def digest(self):\n        return hashlib.sha256(inspect.getsource(self.__code).encode()).hexdigest()\n\n\ndef predict_code(w, xs):\n    return 1 if float(w[:-1] @ xs + w[-1]) > 0 else -1\n\n\nclass NormalWorld:\n    \"\"\"What the display software can do: ask for predictions.\"\"\"\n\n    def __init__(self, secure):\n        self._secure = secure\n\n    def flag(self, x):\n        return self._secure.predict(x)\n\n\nsecure = SecureWorld(w_hat, predict_code)\nnormal = NormalWorld(secure)\npreds = np.array([normal.flag(x) for x in x_new])\nvisible = [a for a in dir(normal) if not a.startswith(\"_\")]\ncheck(f\"[B-isolation]   the normal world's interface offers {visible} and hands back one prediction \"\n      f\"per heartbeat; the model parameters and the recorded heartbeats are not among its values\",\n      visible == [\"flag\"] and set(preds) <= {-1, 1})\nacc = float(np.mean(preds == y_new))\ncheck(f\"[B-isolation]   the predictions it receives flag arrhythmias with accuracy {acc:.2f} \"\n      \"on 200 new heartbeats\", acc > 0.8)\nwith open(OUT_DIR / \"tee_heartbeats.csv\", \"w\") as fh:\n    fh.write(\"interval,qrs_width,r_height,t_height,p_height,st_level,rr_variability,noise,label,prediction\\n\")\n    for xi, yi, pi in zip(x_new, y_new, preds):\n        fh.write(\",\".join(f\"{v:.3f}\" for v in xi) + f\",{yi},{pi}\\n\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-attestation]** The digest of the prediction code, signed with the vendor key, verifies on the client; an altered code that flags fewer arrhythmias changes the digest and fails attestation."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "VENDOR_KEY = b\"vendor-signing-key\"                # known to the chip vendor and the client\n\n\ndef sign(digest):\n    return hmac.new(VENDOR_KEY, digest.encode(), hashlib.sha256).hexdigest()\n\n\ndef verify(digest, signature, expected_digest):\n    return hmac.compare_digest(signature, sign(digest)) and digest == expected_digest\n\n\nexpected = hashlib.sha256(inspect.getsource(predict_code).encode()).hexdigest()\nd_ok = secure.digest()\ncheck(f\"[B-attestation] the client verifies the signed digest {d_ok[:12]}... of the prediction \"\n      \"code before sending heartbeats\", verify(d_ok, sign(d_ok), expected))\n\n\ndef tampered_code(w, xs):\n    return 1 if float(w[:-1] @ xs + w[-1]) > 2.0 else -1   # flags fewer arrhythmias\n\n\ntampered = SecureWorld(w_hat, tampered_code)\nflags_ok = sum(normal.flag(x) == 1 for x in x_new)\nflags_t = sum(NormalWorld(tampered).flag(x) == 1 for x in x_new)\nd_t = tampered.digest()\ncheck(f\"[B-attestation] altered code that flags {flags_t} instead of {flags_ok} heartbeats has the \"\n      f\"digest {d_t[:12]}... and fails attestation\", flags_t < flags_ok and not verify(d_t, sign(d_t), expected))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-reveal]** What the prediction reveals is not protected: on the 20 recorded heartbeats every score has the sign of the label and a large size, while on new heartbeats the scores are smaller and sometimes wrong, so a membership inference attack that only observes predictions and knows the label tells recorded from new heartbeats well above chance."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# the signed score y * score: positive and large when the prediction is right with\n# a large score.  An attacker who holds a heartbeat with its label and sees\n# the score guesses \"recorded\" when the signed score exceeds a threshold.\nsigned = lambda x, y: y * (np.c_[standardize(x), np.ones(len(x))] @ w_hat)\nm_rec, m_new = signed(x_rec, y_rec), signed(x_new, y_new)\nthresholds = np.unique(np.r_[m_rec, m_new])\nbalanced = [0.5 * (np.mean(m_rec > t) + np.mean(m_new <= t)) for t in thresholds]\nmia = float(max(balanced)); t_best = float(thresholds[int(np.argmax(balanced))])\ncheck(f\"[B-reveal]      every recorded heartbeat has a positive signed score (smallest {m_rec.min():.1f}), while \"\n      f\"{int((m_new <= 0).sum())} of the {N_NEW} new ones have a negative signed score and the rest a typical signed score \"\n      f\"of {np.median(m_new):.1f} against {np.median(m_rec):.1f} on recorded ones; guessing 'recorded' above the \"\n      f\"signed score {t_best:.1f} is right {mia:.2f} of the time on a balanced mix, above chance 0.5\",\n      m_rec.min() > 0 and mia > 0.65)\nwith open(OUT_DIR / \"tee_scores.csv\", \"w\") as fh:\n    fh.write(\"recorded,signed_score\\n\")\n    for v in m_rec: fh.write(f\"1,{v:.4f}\\n\")\n    for v in m_new: fh.write(f\"0,{v:.4f}\\n\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[B-sidechannel]** A power trace modeled as the number of one-bits of each operand in the multiplications of w^T x leaks the model parameters: the traces of different heartbeats differ with the feature values, and from them the sign and size of each weight is recovered; constant-time code processes every operand the same way and gives a flat trace."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "def bits(v):\n    \"\"\"Number of one-bits of a value in 8-bit fixed point: the power the multiplier draws.\"\"\"\n    return bin(int(round(abs(v) * 16)) & 0xFF).count(\"1\")\n\n\ndef trace_data_dependent(w, xs):\n    \"\"\"One multiplication per feature; the operand bits leave a trace.\"\"\"\n    return [bits(wk) + bits(xk) for wk, xk in zip(w[:-1], xs)]\n\n\ndef trace_constant_time(w, xs):\n    \"\"\"Every operand masked to a full word before use: the trace is flat.\"\"\"\n    return [8 + 8 for _ in zip(w[:-1], xs)]\n\n\nxs0 = standardize(x_new[:1])[0]\nt_dep, t_ct = trace_data_dependent(w_hat, xs0), trace_constant_time(w_hat, xs0)\ntraces = np.array([trace_data_dependent(w_hat, standardize(x_new[i:i + 1])[0]) for i in range(50)])\nbits_x = np.array([[bits(v) for v in row] for row in standardize(x_new[:50])])\nrecovered = traces - bits_x                           # what remains is the weights' contribution\ncheck(f\"[B-sidechannel] the traces of 50 heartbeats differ with the data (spread \"\n      f\"{list(traces.max(0) - traces.min(0))} per multiplication), and subtracting the operand bits of each \"\n      f\"heartbeat leaves the bits of the weights, {list(recovered[0])}, the same for every heartbeat\",\n      bool(np.all(recovered == [bits(v) for v in w_hat[:-1]])))\ncheck(f\"[B-sidechannel] constant-time code gives the trace {t_ct} for every heartbeat\",\n      len(set(t_ct)) == 1 and trace_constant_time(w_hat, standardize(x_new[1:2])[0]) == t_ct)\nwith open(OUT_DIR / \"tee_trace.csv\", \"w\") as fh:\n    fh.write(\"multiplication,data_dependent,constant_time\\n\")\n    for k, (a, b) in enumerate(zip(t_dep, t_ct), 1):\n        fh.write(f\"{k},{a},{b}\\n\")\n\n# -------------------------------------------------------------- preview\nfig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(12.5, 3.7))\nfor lab, mk, name in ((-1, \"o\", \"normal\"), (1, \"^\", \"arrhythmia\")):\n    m = y_new == lab\n    ax1.scatter(x_new[m, 0], x_new[m, 1], marker=mk, s=18, facecolors=\"none\" if lab == 1 else \"k\",\n                edgecolors=\"k\", label=f\"{name}, {np.mean(preds[m] == lab):.2f} flagged right\")\nax1.set_xlabel(FEATURES[0]); ax1.set_ylabel(FEATURES[1])\nax1.set_title(\"what the normal world sees: predictions\"); ax1.legend(frameon=False, fontsize=7)\nax2.hist(m_rec, bins=20, histtype=\"step\", color=\"k\", label=\"recorded heartbeats\")\nax2.hist(m_new, bins=20, histtype=\"step\", color=\"k\", ls=\"--\", label=\"new heartbeats\")\nax2.set_xlabel(\"signed score $y\\\\,\\\\mathbf{w}^\\\\top \\\\mathbf{x}$\"); ax2.set_ylabel(\"number of heartbeats\")\nax2.set_title(\"what the predictions reveal\"); ax2.legend(frameon=False, fontsize=8)\nk = np.arange(1, 9)\nax3.plot(k, t_dep, \"k-o\", ms=5, label=\"data-dependent code\")\nax3.plot(k, t_ct, \"k--s\", ms=5, mfc=\"none\", label=\"constant-time code\")\nax3.set_xticks(k); ax3.set_xlabel(\"multiplication of $\\\\mathbf{w}^\\\\top \\\\mathbf{x}$\"); ax3.set_ylabel(\"operand one-bits (power)\")\nax3.set_title(\"the power trace side channel\"); ax3.legend(frameon=False, fontsize=8)\nfig.tight_layout()\nfig.savefig(OUT_DIR / \"tee.png\", dpi=110)\n\nn_ok = sum(ok for _, ok in report)\nprint(f\"\\n{n_ok}/{len(report)} checks pass\")\nprint(f\"wrote {OUT_DIR / 'tee_heartbeats.csv'}, {OUT_DIR / 'tee_scores.csv'}, {OUT_DIR / 'tee_trace.csv'}, {OUT_DIR / 'tee.png'}\")\nif n_ok != len(report):\n    raise SystemExit(1)"
  }
 ]
}