localstack-pro-samples/transfer-ftp-s3/test.py at master · localstack-samples/localstack-pro-samples · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import io
import time
import boto3
from ftplib import FTP
import re

EDGE_URL = 'http://localhost:4566'

USERNAME = 'user_123'
BUCKET = 'transfer-files'
S3_DIR = 'mydir'
S3_FILENAME = 'test-file-aws-transfer.txt'
FTP_USER_DEFAULT_PASSWD = '12345'
FILE_CONTENT = b'title "Test" \nfile content!!'


def create_transfer_api():
    transfer_client, s3_client = get_clients()

    print('Creating FTP server in AWS Transfer API')

    role_arn = 'arn:aws:iam::testrole'

    rs = transfer_client.create_server(
        EndpointType='PUBLIC',
        IdentityProviderType='SERVICE_MANAGED',
        Protocols=['FTP']
    )
    time.sleep(1)

    server_id = rs['ServerId']
    match = re.match(r"^s-[a-z]*([0-9]{4,5})$", server_id)
    port = int(match.group(1))

    s3_client.create_bucket(Bucket=BUCKET)

    transfer_client.create_user(
        ServerId=server_id,
        HomeDirectory=BUCKET,
        HomeDirectoryType='PATH',
        Role=role_arn,
        UserName=USERNAME
    )
    return server_id, port


def upload_files(server_id, ftp_port):
    transfer_client, s3_client = get_clients()

    ftp = FTP()
    print('Connecting to AWS Transfer FTP server on local port %s' % ftp_port)
    ftp.connect('localhost', port=ftp_port)

    # connect to FTP server
    result = ftp.login(USERNAME, FTP_USER_DEFAULT_PASSWD)
    assert 'Login successful.' in result

    # upload file to root dir
    print('Uploading file to FTP root directory')
    retry(ftp.storbinary, cmd='STOR %s' % S3_FILENAME, fp=io.BytesIO(FILE_CONTENT))

    # upload file to sub dir
    print('Uploading file to FTP sub-directory')
    ftp.mkd(S3_DIR)
    ftp.cwd(S3_DIR)
    retry(ftp.storbinary, cmd='STOR %s' % S3_FILENAME,fp=io.BytesIO(FILE_CONTENT))

    ftp.quit()


def download_files(server_id, ftp_port):
    transfer_client, s3_client = get_clients()

    print('Downloading files from S3 root and sub-directory')

    rs = s3_client.get_object(Bucket=BUCKET, Key=S3_FILENAME)
    assert rs['Body'].read() == FILE_CONTENT

    rs = s3_client.get_object(Bucket=BUCKET, Key='{}/{}'.format(S3_DIR, S3_FILENAME))
    assert rs['Body'].read() == FILE_CONTENT


def get_clients():
    return boto3.client('transfer', endpoint_url=EDGE_URL), boto3.client('s3', endpoint_url=EDGE_URL)

# copied from localstack, to avoid "connection refused" errors

def retry(function, retries=3, sleep=1.0, sleep_before=0, **kwargs):
    raise_error = None
    if sleep_before > 0:
        time.sleep(sleep_before)
    retries = int(retries)
    for i in range(0, retries + 1):
        try:
            return function(**kwargs)
        except Exception as error:
            raise_error = error
            time.sleep(sleep)
    raise raise_error

def main():
    server_id, ftp_port = create_transfer_api()
    upload_files(server_id, ftp_port)
    download_files(server_id, ftp_port)
    print('Tests succesfully completed.')


if __name__ == '__main__':
    main()