moebiusT7/book-ocr-studio

1"""Authenticated loopback bridge for the user's existing Chrome window."""2import base64,fcntl,json,queue,secrets,threading,time,uuid3from http.server import BaseHTTPRequestHandler,ThreadingHTTPServer4from pathlib import Path5from core import ROOT,read,save6PORT=85087KEY=ROOT/'.chrome-bridge-key'8def key():9    if not KEY.exists():10        try:11            with KEY.open('x') as f:f.write(secrets.token_urlsafe(32))12            KEY.chmod(0o600)13        except FileExistsError:pass14    return KEY.read_text().strip()15 16def call(action,**args):17    import requests18    r=requests.post(f'http://127.0.0.1:{PORT}/rpc',json=dict(action=action,**args),headers={'X-Book-OCR-Key':key()},timeout=75)19    r.raise_for_status();data=r.json()20    if data.get('error'):raise RuntimeError(data['error'])21    return data.get('result')22 23class Mouse:24    def click(self,x,y):call('click',x=x,y=y)25    def move(self,x,y):call('move',x=x,y=y)26class ChromePage:27    def __init__(self):28        info=call('connect');self.viewport_size=dict(width=info['width'],height=info['height']);self.url=info['url'];self._title=info['title']29        self.frames=[self];self.main_frame=self;self.mouse=Mouse()30    def evaluate(self,script,arg=None):return call('evaluate',script=script,arg=arg)31    def bring_to_front(self):pass32    def wait_for_timeout(self,ms):time.sleep(ms/1000)33    def title(self):return self._title34    def screenshot(self,clip=None,animations=None,**kw):35        data=base64.b64decode(call('screenshot',clip=clip))36        if kw.get('path'):Path(kw['path']).write_bytes(data)37        return data38 39def run():40    requests_q=queue.Queue();pending={};lock=threading.Lock()41    class H(BaseHTTPRequestHandler):42        def log_message(self,*a):pass43        def do_GET(self):44            if self.path!='/health':self.send_error(404);return45            self.send_response(200);self.send_header('Content-Type','application/json');self.end_headers();self.wfile.write(b'{"service":"book-ocr-chrome"}')46        def do_POST(self):47            if self.headers.get('X-Book-OCR-Key')!=key():self.send_error(403);return48            if int(self.headers.get('Content-Length',0))>20_000_000:self.send_error(413);return49            try:50                data=json.loads(self.rfile.read(int(self.headers.get('Content-Length',0))))51                if self.path=='/poll':52                    save(ROOT/'chrome-bridge-status.json',dict(connected=True,seen=time.time()))53                    try:result=requests_q.get(timeout=1)54                    except queue.Empty:result=None55                elif self.path=='/reply':56                    with lock:entry=pending.get(data['id'])57                    if entry:entry['data']=data;entry['event'].set()58                    result={}59                elif self.path=='/rpc':60                    if data.get('action') not in {'connect','evaluate','click','move','screenshot','disconnect'}:self.send_error(400);return61                    ident=uuid.uuid4().hex;entry=dict(event=threading.Event())62                    with lock:pending[ident]=entry63                    requests_q.put(dict(id=ident,deadline=time.time()+65,**data))64                    ok=entry['event'].wait(66)65                    with lock:pending.pop(ident,None)66                    result=entry['data'] if ok else dict(error='Cannot connect to the Chrome extension. Check the extension and Kindle tab.')67                else:self.send_error(404);return68                body=json.dumps(result).encode();self.send_response(200);self.send_header('Content-Type','application/json');self.end_headers();self.wfile.write(body)69            except (BrokenPipeError,ConnectionResetError):pass70            except Exception as exc:self.send_error(500,str(exc))71    ThreadingHTTPServer(('127.0.0.1',PORT),H).serve_forever()72if __name__=='__main__':run()73