One-way Functions

import Mathlib.Probability.Distributions.Uniform
import Onewayf.Adversary
import Onewayf.Negligible

-- leandown
-- [meta]
-- title = "One-way functions"
-- group = "Cryptography"
-- [content]

-- # {{definition}}: Inversion probability
--
-- The **inversion probability** of adversary `F` against function `f` at
-- security parameter `n` is the probability that `F` finds a preimage of
-- `f(x)` when `x` is drawn uniformly from $\{0,1\}^n$:
--
-- $$\mathrm{invertProb}(f, F, n) =
--   \sum_{x \in \{0,1\}^n} \frac{1}{2^n}
--   \cdot \Pr_{x' \leftarrow F(f(x))}[f(x') = f(x)]$$
--
-- Note that `F` only needs to find *some* preimage of `f(x)`, not necessarily
-- `x` itself.
--
noncomputable def invertProb
    (f : BitString → BitString)
    (F : BitString → PMF BitString)
    (n : ℕ) : ℝ :=
  (∑ x : FixedBitString n,
    PMF.uniformOfFintype (FixedBitString n) x *
    (∑' x' : BitString, F (f x.toList) x' *
      (if f x' = f x.toList ∧ x'.length = n then 1 else 0)) : ENNReal).toReal

-- # {{definition}}: One-way functions
--
-- A function `f : BitString → BitString` is a **one-way function** if:
-- 1. **Efficiency** — `f` is poly-time computable.
-- 2. **Hardness** — for every poly-time adversary `F`, the inversion
--    probability `invertProb f F n` is negligible in `n`.
--
-- Hardness says that no efficient algorithm can invert `f` on a noticeable
-- fraction of inputs, even when the input is chosen uniformly at random.
--
structure IsOneWayFunction (f : BitString → BitString) : Prop where
  polytime : PolyTimeComputable f
  hard     : ∀ F : BitString → PMF BitString, PolyTimeAdversary F →
               Negligible (invertProb f F)

-- # {{lemma}}: A perfect inverter achieves probability 1
--
-- If `F` always returns `pure x` on input `f(x)` — i.e. it deterministically
-- recovers the exact preimage — then `invertProb f F n = 1` for all `n`.
-- This is the strongest possible inversion guarantee.
--
lemma invertProb_pure_self_eq_one
    (f : BitString → BitString) (F : BitString → PMF BitString)
    (hF : ∀ x : BitString, F (f x) = PMF.pure x)
    (n : ℕ) : invertProb f F n = 1 := by
  simp only [invertProb]
  have inner : ∀ x : FixedBitString n,
      (∑' x' : BitString, F (f x.toList) x' *
        (if f x' = f x.toList ∧ x'.length = n then (1 : ENNReal) else 0)) = 1 := by
    intro x
    rw [hF]
    have h : ∀ x' : BitString,
        (PMF.pure x.toList) x' * (if f x' = f x.toList ∧ x'.length = n then (1 : ENNReal) else 0) =
        (PMF.pure x.toList) x' := by
      intro x'
      simp only [PMF.pure_apply]
      by_cases heq : x' = x.toList
      · subst heq; simp [x.toList_length]
      · simp [heq]
    simp_rw [h]
    exact PMF.tsum_coe _
  simp_rw [inner, mul_one]
  have outer : (∑ x : FixedBitString n,
      PMF.uniformOfFintype (FixedBitString n) x) = 1 := by
    have := (PMF.uniformOfFintype (FixedBitString n)).tsum_coe
    rw [tsum_fintype (L := .unconditional _)] at this; exact this
  simp

-- # {{lemma}}: A preimage inverter achieves probability 1
--
-- A weaker sufficient condition: if for every `x : FixedBitString n`, `F(f(x.toList))` is a
-- point mass at some `x'` satisfying `f(x') = f(x.toList)` and `x'.length = n`,
-- then `invertProb f F n = 1`.
-- This handles cases where the recovered preimage may differ from `x` — for
-- example, when `f` is not injective or when bit normalization means the
-- canonical preimage differs from the sampled one.
--
lemma invertProb_pure_preimage_eq_one
    (f : BitString → BitString) (F : BitString → PMF BitString)
    (hF : ∀ (n : ℕ) (x : FixedBitString n),
            ∃ x' : BitString, F (f x.toList) = PMF.pure x' ∧ f x' = f x.toList ∧ x'.length = n)
    (n : ℕ) : invertProb f F n = 1 := by
  simp only [invertProb]
  have inner : ∀ x : FixedBitString n,
      (∑' x' : BitString, F (f x.toList) x' *
        (if f x' = f x.toList ∧ x'.length = n then (1 : ENNReal) else 0)) = 1 := by
    intro x
    obtain ⟨x', hFx, hfx', hlen⟩ := hF n x
    rw [hFx]
    have h : ∀ x'' : BitString,
        (PMF.pure x') x'' * (if f x'' = f x.toList ∧ x''.length = n then (1 : ENNReal) else 0) =
        (PMF.pure x') x'' := by
      intro x''
      simp only [PMF.pure_apply]
      by_cases heq : x'' = x'
      · subst heq; simp [hfx', hlen]
      · simp [heq]
    simp_rw [h]
    exact PMF.tsum_coe _
  simp_rw [inner, mul_one]
  have outer : (∑ x : FixedBitString n,
      PMF.uniformOfFintype (FixedBitString n) x) = 1 := by
    have := (PMF.uniformOfFintype (FixedBitString n)).tsum_coe
    rw [tsum_fintype (L := .unconditional _)] at this; exact this
  simp