"""A/B transaction model; does not claim physical flash atomicity."""
from hashlib import sha256


class Update:
    def __init__(self):
        self.active = "A"
        self.ready = None
        self.candidate = None
        self.tried = False

    def download(self, image, expected_digest):
        if sha256(image).hexdigest() != expected_digest:
            return False
        self.ready = "B" if self.active == "A" else "A"
        return True

    def select(self):
        if self.ready is None:
            raise ValueError("verified candidate required")
        self.candidate = self.ready
        self.ready = None
        self.tried = False

    def boot(self):
        if self.candidate is None:
            return self.active
        if self.tried:
            self.candidate = None
            return self.active
        self.tried = True
        return self.candidate

    def confirm(self):
        if self.candidate is None or not self.tried:
            raise ValueError("candidate must be running")
        self.active = self.candidate
        self.candidate = None
        self.tried = False


def main():
    image = b"synthetic candidate firmware version 2"
    digest = sha256(image).hexdigest()
    broken = Update()
    assert not broken.download(image + b"damage", digest)
    assert broken.boot() == "A"
    before_select = Update()
    assert before_select.download(image, digest)
    assert before_select.boot() == "A"
    failed_trial = Update()
    assert failed_trial.download(image, digest)
    failed_trial.select()
    assert failed_trial.boot() == "B"
    assert failed_trial.boot() == "A"
    successful = Update()
    assert successful.download(image, digest)
    successful.select()
    assert successful.boot() == "B"
    successful.confirm()
    assert successful.boot() == "B"
    print("corruption -> A; unselected -> A; failed trial -> A; confirmed -> B")
    print("PASS")


if __name__ == "__main__":
    main()