KnowHow

技術的なメモを中心にまとめます。
検索にて調べることができます。

asyncioを用いてサーバ側のコマンド結果をクライアントで取得するコードを書いてみた。

登録日 :2024/09/18 05:48
カテゴリ :Python基礎

サーバ側で適当なコマンドの実行結果をクライアント側で受け取る簡易的なAPIサーバを構築したい。
できるだけ、サードパーティのライブラリを用いず、デフォルトでインストールずみのpythonライブラリで構築できないかを考えた。
socketだけだとちょっと難しそうだった。asyncioを用いると、やりたいことに近いコードが書けそうだったので、サンプルコードを作成してみた。

サーバ側

クライアント側からのキックを待つ。クライアントからの入力をname変数に受け取って、サーバ側のコマンドに使っている。ただし、以下の書き方では、コマンドインジェクション攻撃が可能なので、修正が必要であることは間違いない。。とりあえず、テストまではできたので忘備録としてまとめた。
server.py

import asyncio
import asyncio.subprocess
import collections


class CounterServer(object):
    def __init__(self):
        self.counter = collections.Counter()
        self.lock = asyncio.Lock()
        self.base_command = 'ls -l | grep '

    async def handle_echo(self, reader, writer):
        data = await reader.read()
        name = data.decode()

        # with await self.lock:
        async with self.lock:
            if self.counter[name] > 10:
                writer.write(b'-1')
                self.counter[name] = 0
            else:
                writer.write(str(self.counter[name]).encode())
                self.counter[name] += 1
        await writer.drain()
        writer.close()

    async def run_command(self, reader, writer):
        data = await reader.read()
        name = data.decode()
        print(name)

        async with self.lock:
            proc = await asyncio.create_subprocess_shell(
                self.base_command + name, stdout=asyncio.subprocess.PIPE
            )
            stdout, stderr = await proc.communicate()
            res = str(stdout.decode())
            writer.write(res.encode())
            exitcode = await proc.wait()
            if exitcode == 1:
                writer.write(str('error:' + name).encode())
            else:
                print(exitcode)
        await writer.drain()
        writer.close()


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    counter_sever = CounterServer()
    # coro = asyncio.start_server(counter_sever.handle_echo,
    #                             '127.0.0.1', 8888, loop=loop)
    coro = asyncio.start_server(counter_sever.run_command,
                                '127.0.0.1', 8888, loop=loop)

    server = loop.run_until_complete(coro)
    print('server {}'.format(server.sockets[0].getsockname()))
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass

    server.close()
    loop.run_until_complete(server.wait_closed())
    loop.close()

クライアント側

クライアントがキックして、サーバ側のコマンド入力結果を取得する。

import asyncio


class AwaitableClass(object):
    def __init__(self, name, _loop):
        self.name = name
        self.loop = _loop

    def __await__(self):
        async def request_server():
            reader, writer = await asyncio.open_connection(
                '127.0.0.1', 8888)
            writer.write(self.name.encode())
            writer.write_eof()
            data = await reader.read()
            data = data.decode()
            return data
        return request_server().__await__()


async def main(name, _loop):
    print('chunk reader')
    result = await AwaitableClass(name, _loop)
    print(result)


if __name__ == '__main__':
    loop = asyncio.get_event_loop()
    loop.run_until_complete(asyncio.wait([
        main('server', loop),
        main('task2', loop),
        main('task3', loop),
    ]))
    loop.close()

実行結果

サーバ側

$ python server.py
server ('127.0.0.1', 8888)
task3 <ークライアントからのリクエスト
task2 <ークライアントからのリクエスト
server <ークライアントからのリクエスト
0

クライアント側

$ python client.py
chunk reader
chunk reader
chunk reader
error:task3
error:task2
-rw-r--r--@ 1 staff  staff  1999  9 18 05:46 server.py