Avinashnalla7 commited on
Commit
47956b3
·
1 Parent(s): 987c523

Add SFTP storage helper

Browse files
Files changed (1) hide show
  1. backend/sftp_store.py +84 -0
backend/sftp_store.py ADDED
@@ -0,0 +1,84 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ def _rjoin(*parts):
4
+ out = []
5
+ for p in parts:
6
+ if not p:
7
+ continue
8
+ p = str(p).strip("/")
9
+ if p:
10
+ out.append(p)
11
+ return "/".join(out)
12
+
13
+
14
+ import base64
15
+ import os
16
+ from pathlib import Path
17
+ from typing import Optional, Tuple
18
+
19
+ import paramiko
20
+
21
+
22
+ def _sftp_client() -> Tuple[paramiko.SFTPClient, paramiko.Transport]:
23
+ host = (os.environ.get("SFTP_HOST") or "").strip()
24
+ user = (os.environ.get("SFTP_USER") or "").strip()
25
+ if not host or not user:
26
+ raise RuntimeError("Missing SFTP_HOST or SFTP_USER")
27
+
28
+ port = int((os.environ.get("SFTP_PORT") or "22").strip())
29
+ password = os.environ.get("SFTP_PASS")
30
+ key_b64 = os.environ.get("SFTP_KEY_B64")
31
+ key_pass = os.environ.get("SFTP_KEY_PASSPHRASE")
32
+
33
+ t = paramiko.Transport((host, port))
34
+ if key_b64:
35
+ key_bytes = base64.b64decode(key_b64)
36
+ pkey = paramiko.RSAKey.from_private_key_file(_write_tmp_key(key_bytes), password=key_pass)
37
+ t.connect(username=user, pkey=pkey)
38
+ else:
39
+ if not password:
40
+ raise RuntimeError("Missing SFTP_PASS (or provide SFTP_KEY_B64)")
41
+ t.connect(username=user, password=password)
42
+
43
+ return paramiko.SFTPClient.from_transport(t), t
44
+
45
+
46
+ def _write_tmp_key(key_bytes: bytes) -> str:
47
+ # HF container FS is writable for tmp; keep it simple.
48
+ p = Path("/tmp/sftp_key.pem")
49
+ p.write_bytes(key_bytes)
50
+ os.chmod(p, 0o600)
51
+ return str(p)
52
+
53
+
54
+ def _mkdir_p(sftp: paramiko.SFTPClient, remote_dir: str) -> None:
55
+ parts = remote_dir.strip("/").split("/")
56
+ cur = ""
57
+ for part in parts:
58
+ cur += "/" + part
59
+ try:
60
+ sftp.stat(cur)
61
+ except FileNotFoundError:
62
+ sftp.mkdir(cur)
63
+
64
+
65
+ def store_to_sftp(pdf_id: str, template_id: str, cfg_json_bytes: bytes, pdf_bytes: bytes, pdf_name: str) -> str:
66
+ base = (os.environ.get("SFTP_BASE_DIR") or "/").rstrip("/")
67
+ remote_dir = f"{base}/{template_id}/{pdf_id}"
68
+ remote_cfg = f"{remote_dir}/trainer_config_{pdf_id}__{template_id}.json"
69
+ remote_pdf = f"{remote_dir}/{pdf_name or (pdf_id + '.pdf')}"
70
+
71
+ sftp, transport = _sftp_client()
72
+ try:
73
+ _mkdir_p(sftp, remote_dir)
74
+ with sftp.open(remote_cfg, "wb") as f:
75
+ f.write(cfg_json_bytes)
76
+ with sftp.open(remote_pdf, "wb") as f:
77
+ f.write(pdf_bytes)
78
+ finally:
79
+ try:
80
+ sftp.close()
81
+ finally:
82
+ transport.close()
83
+
84
+ return remote_dir