42 lines
2.5 KiB
Python
42 lines
2.5 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: GPL-3.0-or-later
|
|
"""Network-free streaming receiver for one Phase-1.0DM result."""
|
|
from __future__ import annotations
|
|
from dataclasses import dataclass
|
|
import hashlib,os
|
|
from pathlib import Path
|
|
from phase10dm_snapshot_protocol import BEGIN,END,FRAME,MAGIC,MAX_BYTES,ProtocolError,_frame
|
|
|
|
@dataclass(frozen=True)
|
|
class ReceivedSnapshot:
|
|
path:Path;size:int;sha256:str;device:int;inode:int;mtime:int
|
|
|
|
class SnapshotReceiver:
|
|
def __init__(self,path:Path):
|
|
if not isinstance(path,Path) or path.exists():raise ProtocolError("exclusive output required")
|
|
self.path=path;self.stream=path.open("xb");self.buffer=bytearray();self.state="HEADER";self.expected=None;self.written=0;self.digest=hashlib.sha256();self.begin=None
|
|
def feed(self,chunk:bytes):
|
|
if self.state not in {"HEADER","PAYLOAD","TERMINAL"} or not isinstance(chunk,bytes):raise ProtocolError("receiver state")
|
|
self.buffer.extend(chunk)
|
|
if self.state=="HEADER" and len(self.buffer)>=64:
|
|
self.begin=_frame(bytes(self.buffer[:64]));del self.buffer[:64];kind,status,saved,size,sent,*_=self.begin[2:]
|
|
if kind!=BEGIN or status or saved or sent or not 0<size<=MAX_BYTES:self.abort();raise ProtocolError("begin invalid")
|
|
self.expected=size;self.state="PAYLOAD"
|
|
if self.state=="PAYLOAD":
|
|
take=min(len(self.buffer),self.expected-self.written)
|
|
if take:
|
|
data=bytes(self.buffer[:take]);del self.buffer[:take];self.stream.write(data);self.digest.update(data);self.written+=take
|
|
if self.written==self.expected:self.state="TERMINAL"
|
|
if self.state=="TERMINAL" and len(self.buffer)>64:self.abort();raise ProtocolError("trailing bytes")
|
|
def finish(self)->ReceivedSnapshot:
|
|
if self.state!="TERMINAL" or len(self.buffer)!=64 or self.begin is None:self.abort();raise ProtocolError("stream incomplete")
|
|
end=_frame(bytes(self.buffer));kind,status,saved,size,sent,dev,ino,mtime=end[2:];_,_,_,bsize,_,bdev,bino,bmtime=self.begin[2:]
|
|
if kind!=END or status or saved or size!=bsize or sent!=bsize or (dev,ino,mtime)!=(bdev,bino,bmtime):self.abort();raise ProtocolError("terminal invalid")
|
|
self.stream.flush();os.fsync(self.stream.fileno());self.stream.close();self.state="SEALED"
|
|
with self.path.open("rb") as check:
|
|
if check.read(16)!=b"SQLite format 3\x00":raise ProtocolError("not sqlite")
|
|
return ReceivedSnapshot(self.path,self.written,self.digest.hexdigest(),dev,ino,mtime)
|
|
def abort(self):
|
|
if not self.stream.closed:self.stream.flush();self.stream.close()
|
|
self.state="ABORTED"
|