import asyncio
import asyncssh
from fastapi import FastAPI
from fastapi import Query
from fastapi_range_response import FileRangeResponse, AsyncsshRangeResponse

app = FastAPI()


@app.get('/local_file')
async def download_local_file(full_path: str = Query()):
    return FileRangeResponse(full_path, media_type='application/octet-stream')


@app.get('/download_sftp_file')
async def download_sftp_file(full_path: str = Query()):
    ssh_client = await asyncio.wait_for(
        asyncssh.connect(host='xxx', port=22, username='xxx', password='xxxx', known_hosts=None),
        timeout=10
    )
    return AsyncsshRangeResponse(full_path, ssh_client, done_close=True)


if __name__ == '__main__':
    import uvicorn

    uvicorn.run(app)