feat: added traps command

This commit is contained in:
Krzysztof Rudnicki 2024-06-24 12:05:31 +02:00
parent bd3bb4a4cb
commit 14d6807119
2 changed files with 36 additions and 51 deletions

84
main.py
View File

@ -4,6 +4,7 @@ import websockets
import requests import requests
import base64 import base64
import json import json
from rule34Py import rule34Py
from datetime import datetime, time, timedelta from datetime import datetime, time, timedelta
@ -15,15 +16,7 @@ GROUP_ID = os.getenv('GROUP_ID', '')
GROUP_ID_SEND = os.getenv('GROUP_ID_SEND', '') GROUP_ID_SEND = os.getenv('GROUP_ID_SEND', '')
CAT_API = os.getenv('CAT_API', '') CAT_API = os.getenv('CAT_API', '')
def fetch_and_download_dog_image(): def download_image(image_url):
# Send request to The Cat API
response = requests.get("https://dog.ceo/api/breeds/image/random")
response.raise_for_status() # Ensure the request was successful
# Parse the response JSON to get the image URL
data = response.json()
image_url = data['message']
# Download the image # Download the image
image_response = requests.get(image_url) image_response = requests.get(image_url)
image_response.raise_for_status() # Ensure the request was successful image_response.raise_for_status() # Ensure the request was successful
@ -41,38 +34,30 @@ def fetch_and_download_dog_image():
os.remove(image_filename) os.remove(image_filename)
return base64_encoded_data return base64_encoded_data
def fetch_and_download_image(api_url, image_url_key):
def fetch_and_download_cat_image(): # Send request to the API
# Send request to The Cat API response = requests.get(api_url)
response = requests.get("https://api.thecatapi.com/v1/images/search")
response.raise_for_status() # Ensure the request was successful response.raise_for_status() # Ensure the request was successful
# Parse the response JSON to get the image URL # Parse the response JSON to get the image URL
print("response: ", response.json)
data = response.json() data = response.json()
image_url = data[0]['url']
# Download the image # Retrieve the image URL based on the provided key
image_response = requests.get(image_url) if isinstance(image_url_key, list):
image_response.raise_for_status() # Ensure the request was successful image_url = data
for key in image_url_key:
image_url = image_url[key]
else:
image_url = data[image_url_key]
# Extract the image filename from the URL return download_image(image_url)
image_filename = image_url.split("/")[-1]
with open(image_filename, 'wb') as image_file: def send_image(base64_attachments, recipients=GROUP_ID_SEND):
image_file.write(image_response.content)
# Convert the image to base64 encoded data
with open(image_filename, 'rb') as image_file:
base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8')
os.remove(image_filename)
return base64_encoded_data
def send_cat():
data = { data = {
"base64_attachments": [fetch_and_download_cat_image()], "base64_attachments": [base64_attachments],
"number": PHONE_NUMBER, "number": PHONE_NUMBER,
"recipients": [GROUP_ID_SEND] "recipients": [recipients]
} }
response = requests.post(SEND_URL, json=data) response = requests.post(SEND_URL, json=data)
if response.status_code == 200: if response.status_code == 200:
@ -81,24 +66,12 @@ def send_cat():
print(f"Request failed with status code: {response.status_code}") print(f"Request failed with status code: {response.status_code}")
print(response.text) print(response.text)
def send_dog(): def send_message(message_content, recipients=GROUP_ID_SEND):
data = {
"base64_attachments": [fetch_and_download_dog_image()],
"number": PHONE_NUMBER,
"recipients": [GROUP_ID_SEND]
}
response = requests.post(SEND_URL, json=data)
if response.status_code == 200:
print("Request was successful.")
else:
print(f"Request failed with status code: {response.status_code}")
print(response.text)
def send_message(message_content):
data = { data = {
"message": message_content, "message": message_content,
"number": PHONE_NUMBER, "number": PHONE_NUMBER,
"recipients": [GROUP_ID_SEND] #"recipients": [GROUP_ID_SEND]
"recipients": [recipients]
} }
response = requests.post(SEND_URL, json=data) response = requests.post(SEND_URL, json=data)
if response.status_code == 200: if response.status_code == 200:
@ -137,10 +110,12 @@ def update_string_count(string, mapping):
command_map = { command_map = {
"!kot": send_cat, "!kot": lambda recipient: send_image(fetch_and_download_image("https://api.thecatapi.com/v1/images/search", [0, 'url']), recipient),
"!pies": send_dog "!pies": lambda recipient: send_image(fetch_and_download_image("https://dog.ceo/api/breeds/image/random", 'message'), recipient),
"!traps": lambda recipient: send_image(download_image(((rule34Py().random_post(["trap"])).sample)), recipient)
} }
def update_string_count(string, mapping): def update_string_count(string, mapping):
if string in mapping: if string in mapping:
mapping[string] += 1 mapping[string] += 1
@ -181,12 +156,21 @@ async def listen_to_server(queue):
try: try:
async for message in websocket: async for message in websocket:
message_content = extract_message_content(message) message_content = extract_message_content(message)
print(message_content)
if message_group_id(message_content) == GROUP_ID: if message_group_id(message_content) == GROUP_ID:
await count_messages(message_content, queue) await count_messages(message_content, queue)
message_value = message_message(message_content) message_value = message_message(message_content)
if message_value in command_map: if message_value in command_map:
# Call the corresponding function # Call the corresponding function
command_map[message_value]() command_map[message_value](GROUP_ID_SEND)
else:
print("Unknown command")
elif message_content.get('destinationNumber', {}) == PHONE_NUMBER:
await count_messages(message_content, queue)
message_value = message_message(message_content)
if message_value in command_map:
# Call the corresponding function
command_map[message_value](PHONE_NUMBER)
else: else:
print("Unknown command") print("Unknown command")
except websockets.ConnectionClosed as e: except websockets.ConnectionClosed as e:

View File

@ -4,4 +4,5 @@ flask
uvicorn uvicorn
asyncio asyncio
websockets websockets
requests requests
rule34py