library_receiver.py 9.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147
  1. """Forced-command SSH receiver: fixed registry, source validation, no arbitrary commands."""
  2. import sys
  3. sys.dont_write_bytecode=True
  4. import argparse
  5. import base64
  6. import io
  7. import json
  8. import os
  9. from pathlib import Path, PurePosixPath
  10. import re
  11. import subprocess
  12. import tempfile
  13. import zipfile
  14. import voc_db as v
  15. import library_channel as channel
  16. import agent_batch
  17. import update_library
  18. def hex_sha(value):
  19. v.require(isinstance(value,str) and re.fullmatch('[a-f0-9]{64}',value),'INVALID_FINGERPRINT');return value
  20. def serve_config(path):
  21. data=json.loads(channel.private_file(path).read_text(encoding='utf-8-sig'))
  22. v.require(isinstance(data,dict) and set(data)=={'schema','registry','intake','canPublish'} and data['schema']==channel.VERSION and type(data['canPublish']) is bool,'INVALID_RECEIVER_CONFIG')
  23. for name in ('registry','intake'):
  24. v.require(isinstance(data[name],str) and Path(data[name]).is_absolute(),'ABSOLUTE_RECEIVER_PATH_REQUIRED')
  25. p=Path(data[name]).absolute()
  26. v.require(p.is_absolute() and p.is_dir() and not any(x.is_symlink() for x in (p,*p.parents)),'INVALID_RECEIVER_DIRECTORY')
  27. if os.name!='nt':v.require(p.stat().st_uid==os.getuid() and p.stat().st_mode&0o077==0,'RECEIVER_DIRECTORY_NOT_PRIVATE')
  28. registry,intake=Path(data['registry']).resolve(),Path(data['intake']).resolve()
  29. v.require(registry!=intake and not registry.is_relative_to(intake) and not intake.is_relative_to(registry),'SEPARATE_INTAKE_REQUIRED')
  30. return data
  31. def read(config,request):
  32. v.require(set(request)=={'schema','action','command','params'} and request['command'] in channel.READS,'READ_COMMAND_REQUIRED')
  33. p=request['params'];v.require(isinstance(p,dict) and not set(p)-{'generation','table','profile','entity','id','input','explain'},'INVALID_READ_PARAMS')
  34. argv=[sys.executable,'-B','-X','utf8',str(Path(__file__).with_name('voc_db.py')),request['command'],'--registry',config['registry']]
  35. for k in ('generation','table','entity','id'):
  36. if k in p:
  37. v.require(isinstance(p[k],str) and not p[k].startswith('-') and len(p[k])<=1000,'INVALID_READ_VALUE');argv+=['--'+k,p[k]]
  38. for k in ('profile','explain'):
  39. if k in p:
  40. v.require(type(p[k]) is bool,'INVALID_READ_VALUE')
  41. if p[k]:argv+=['--'+k]
  42. if 'input' in p:argv+=['--input','-']
  43. env=dict(os.environ)
  44. for key in ('KSF_LIBRARY_CHANNEL','KSF_LIBRARY_REGISTRY','KSF_VOC_DB','PYTHONPATH','PYTHONSTARTUP'):env.pop(key,None)
  45. run=subprocess.run(argv,input=v.canonical(p.get('input')).encode('utf-8'),stdout=subprocess.PIPE,stderr=subprocess.PIPE,env=env,timeout=120,check=False)
  46. result=json.loads((run.stdout or run.stderr).decode('utf-8'));v.require(isinstance(result,dict) and type(result.get('ok')) is bool,'INVALID_READ_RESULT')
  47. if result['ok'] and request['command']=='contract':
  48. result['publicationBinding']={'mode':'ssh-channel','newGenerationsSupported':config['canPublish'],
  49. 'remotePublisherEnabled':config['canPublish'],'currentGeneration':result['generation'],'taskAuthorizationStillRequired':True,
  50. 'note':'Agent经受限SSH通道查询同一registry;新批次在后端重建校验后发布,不共享SQLite副本'}
  51. result['dataAccessMode']='ssh-channel';return result
  52. def unpack(raw,root):
  53. with zipfile.ZipFile(io.BytesIO(raw)) as z:
  54. entries=z.infolist();names=set()
  55. v.require(0<len(entries)<=2001 and sum(e.file_size for e in entries)<=channel.MAX_EXPANDED,'INVALID_ARCHIVE_SIZE')
  56. for info in entries:
  57. name=info.filename;parts=PurePosixPath(name).parts
  58. v.require(name not in names and not info.is_dir() and not name.startswith('/') and '\\' not in name and ':' not in name and '..' not in parts and
  59. (name=='batch.json' or name.startswith('raw/')) and (info.external_attr>>16)&0o170000 in (0,0o100000),'UNSAFE_ARCHIVE_MEMBER')
  60. names.add(name)
  61. v.require('batch.json' in names,'BATCH_REQUIRED')
  62. for info in entries:
  63. dest=root/info.filename;dest.parent.mkdir(parents=True,exist_ok=True,mode=0o700)
  64. with dest.open('xb') as f:f.write(z.read(info))
  65. os.chmod(dest,0o600)
  66. path,manifest,batch,rows,assets=__import__('structured_batch').load(root/'batch.json')
  67. v.require(names=={'batch.json',*assets},'UNREFERENCED_ARCHIVE_FILES')
  68. return path
  69. def current(config):
  70. pointer=json.loads((Path(config['registry'])/'current.json').read_text(encoding='utf-8'))
  71. return pointer
  72. def status(config,request):
  73. v.require(set(request)=={'schema','action','params'} and isinstance(request['params'],dict) and set(request['params'])=={'archiveSha256'},'INVALID_STATUS_PARAMS')
  74. digest=hex_sha(request['params']['archiveSha256']);folder=Path(config['intake'])/digest
  75. pointer=current(config);result={'ok':True,'readOnly':True,'databaseWrites':0,'archiveSha256':digest,'currentGeneration':pointer['generation'],'publicationFound':False}
  76. staged=folder/'staging.json'
  77. if staged.is_file():
  78. meta=json.loads(staged.read_text(encoding='utf-8'))
  79. entry=next((x for x in pointer.get('history',[]) if x.get('incomingGeneration')==meta['generation']),None)
  80. if entry:result.update(publicationFound=True,publishedGeneration=entry['generation'],delta=entry.get('delta',{}))
  81. return result
  82. def intake(config,request):
  83. v.require(config['canPublish'],'PUBLICATION_DISABLED')
  84. v.require(set(request)=={'schema','action','params','archiveBase64'},'INVALID_INTAKE_REQUEST')
  85. p=request['params'];v.require(isinstance(p,dict) and set(p)=={'archiveSha256','expectedGeneration','label','apply'} and type(p['apply']) is bool,'INVALID_INTAKE_PARAMS')
  86. digest=hex_sha(p['archiveSha256']);expected=hex_sha(p['expectedGeneration'])
  87. v.require(isinstance(p['label'],str) and 0<len(p['label'])<=200,'INVALID_LABEL')
  88. v.require(current(config)['generation']==expected,'LIBRARY_GENERATION_CHANGED')
  89. v.require(isinstance(request['archiveBase64'],str),'INVALID_ARCHIVE_ENCODING')
  90. raw=base64.b64decode(request['archiveBase64'],validate=True)
  91. v.require(len(raw)<=channel.MAX_ARCHIVE and v.sha(raw)==digest,'ARCHIVE_HASH_MISMATCH')
  92. root=Path(config['intake']);folder=root/digest
  93. if not p['apply']:
  94. with tempfile.TemporaryDirectory(prefix='.validate-',dir=root) as td:
  95. batch=unpack(raw,Path(td));plan=agent_batch.build_batch(batch,Path(td)/'staging.sqlite')
  96. return {**plan,'archiveSha256':digest,'temporaryFilesUsed':True,'persistentFilesWritten':False,'remotePublished':False,'currentGeneration':expected}
  97. if not folder.exists():
  98. with tempfile.TemporaryDirectory(prefix='.receive-',dir=root) as td:
  99. td=Path(td);unpack(raw,td)
  100. (td/'transport.zip').write_bytes(raw);os.chmod(td/'transport.zip',0o600)
  101. # Preserve validated sources before constructing/publishing; retries reuse identical bytes.
  102. td.rename(folder)
  103. td.mkdir() # TemporaryDirectory removes only its empty original path.
  104. else:
  105. v.require(not folder.is_symlink() and (folder/'transport.zip').is_file() and agent_batch.digest(folder/'transport.zip')==digest,'INTAKE_ID_COLLISION')
  106. stage=folder/'staging.sqlite';receipt=folder/'staging.json'
  107. if stage.exists():
  108. v.require(receipt.is_file(),'STAGING_RECEIPT_REQUIRED')
  109. built=json.loads(receipt.read_text(encoding='utf-8'))
  110. v.require(agent_batch.digest(stage)==built['databaseSha256'],'STAGING_HASH_CHANGED')
  111. else:
  112. built=agent_batch.build_batch(folder/'batch.json',stage,True)
  113. with receipt.open('x',encoding='utf-8') as f:f.write(v.canonical(built)+'\n')
  114. os.chmod(receipt,0o600)
  115. update_library.verified(stage)
  116. published=update_library.publish(stage,config['registry'],True,p['label'],expected_generation=expected)
  117. return {**published,'archiveSha256':digest,'remotePublished':True,'sourceArchivePreserved':True,'productionImported':True,
  118. 'databaseWrites':'new-generation-only','readOnly':False,'activeLibraryModified':bool(published.get('applied'))}
  119. def dispatch(config,request):
  120. v.require(isinstance(request,dict) and request.get('schema')==channel.VERSION,'INVALID_PROTOCOL')
  121. action=request.get('action')
  122. if action=='read':return read(config,request)
  123. if action=='status':return status(config,request)
  124. if action=='intake':return intake(config,request)
  125. raise v.LibraryError('ACTION_NOT_ALLOWED')
  126. def main():
  127. p=argparse.ArgumentParser(description=__doc__);p.add_argument('--config',required=True);a=p.parse_args()
  128. applying=False
  129. try:
  130. v.require(os.environ.get('SSH_ORIGINAL_COMMAND') in (None,'','ksf-library-v1'),'SSH_COMMAND_NOT_ALLOWED')
  131. config=serve_config(a.config);wire=sys.stdin.buffer.read(channel.MAX_WIRE+1);v.require(len(wire)<=channel.MAX_WIRE,'CHANNEL_REQUEST_TOO_LARGE')
  132. request=json.loads(wire.decode('utf-8'));applying=request.get('action')=='intake' and request.get('params',{}).get('apply') is True
  133. result=dispatch(config,request);print(v.canonical(result));return 0 if result.get('ok') else 1
  134. except Exception as e:
  135. code=str(e) if isinstance(e,v.LibraryError) else 'RECEIVER_OPERATION_FAILED'
  136. print(v.canonical({'ok':False,'error':{'code':code},'readOnly':not applying,'databaseWrites':None if applying else 0,'publicationOutcome':'not-confirmed'}));return 1
  137. if __name__=='__main__':sys.exit(main())