281 lines
9.2 KiB
Python
281 lines
9.2 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Cron Tool - инструмент для управления задачами пользователя.
|
||
|
||
Позволяет создавать, планировать и выполнять периодические задачи.
|
||
"""
|
||
|
||
import logging
|
||
import sqlite3
|
||
import json
|
||
from pathlib import Path
|
||
from datetime import datetime, timedelta
|
||
from typing import List, Dict, Any, Optional, Callable
|
||
from dataclasses import dataclass, field
|
||
|
||
from bot.tools import BaseTool, ToolResult, register_tool
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class CronJob:
|
||
"""Задача cron."""
|
||
id: Optional[int]
|
||
name: str
|
||
command: str
|
||
schedule: str # cron format: "*/5 * * * *" или "daily", "hourly"
|
||
enabled: bool = True
|
||
last_run: Optional[datetime] = None
|
||
next_run: Optional[datetime] = None
|
||
created_at: datetime = field(default_factory=datetime.now)
|
||
|
||
|
||
class CronTool(BaseTool):
|
||
"""Инструмент для управления задачами пользователя."""
|
||
|
||
name = "cron_manager"
|
||
description = "Управление периодическими задачами пользователя. Создание, планирование и выполнение задач по расписанию."
|
||
category = "automation"
|
||
|
||
def __init__(self, db_path: str = None):
|
||
self.db_path = Path(db_path) if db_path else Path(__file__).parent.parent.parent / "cron.db"
|
||
self._jobs: Dict[int, CronJob] = {}
|
||
self._init_db()
|
||
|
||
def _init_db(self):
|
||
"""Инициализировать БД."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute('''
|
||
CREATE TABLE IF NOT EXISTS cron_jobs (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
name TEXT NOT NULL,
|
||
command TEXT NOT NULL,
|
||
schedule TEXT NOT NULL,
|
||
enabled INTEGER DEFAULT 1,
|
||
last_run DATETIME,
|
||
next_run DATETIME,
|
||
created_at DATETIME DEFAULT CURRENT_TIMESTAMP
|
||
)
|
||
''')
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
def _parse_schedule(self, schedule: str) -> Optional[datetime]:
|
||
"""
|
||
Распарсить расписание и вернуть следующее время выполнения.
|
||
|
||
Поддерживает:
|
||
- "*/N * * * *" - каждые N минут
|
||
- "@hourly" - каждый час
|
||
- "@daily" - каждый день
|
||
- "@weekly" - каждую неделю
|
||
"""
|
||
now = datetime.now()
|
||
|
||
if schedule.startswith('*/'):
|
||
# Каждые N минут
|
||
try:
|
||
minutes = int(schedule.split()[0][2:])
|
||
return now + timedelta(minutes=minutes)
|
||
except (ValueError, IndexError):
|
||
return None
|
||
|
||
elif schedule == '@hourly':
|
||
return now + timedelta(hours=1)
|
||
|
||
elif schedule == '@daily':
|
||
return now + timedelta(days=1)
|
||
|
||
elif schedule == '@weekly':
|
||
return now + timedelta(weeks=1)
|
||
|
||
return None
|
||
|
||
async def add_job(self, name: str, command: str, schedule: str) -> ToolResult:
|
||
"""Добавить задачу."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
|
||
try:
|
||
next_run = self._parse_schedule(schedule)
|
||
next_run_str = next_run.strftime('%Y-%m-%d %H:%M:%S') if next_run else None
|
||
|
||
c.execute('''
|
||
INSERT INTO cron_jobs (name, command, schedule, next_run)
|
||
VALUES (?, ?, ?, ?)
|
||
''', (name, command, schedule, next_run_str))
|
||
|
||
job_id = c.lastrowid
|
||
conn.commit()
|
||
|
||
self._jobs[job_id] = CronJob(
|
||
id=job_id,
|
||
name=name,
|
||
command=command,
|
||
schedule=schedule,
|
||
next_run=next_run
|
||
)
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data={'id': job_id, 'name': name, 'schedule': schedule, 'next_run': next_run_str},
|
||
metadata={'status': 'added'}
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.exception(f"Ошибка добавления задачи: {e}")
|
||
return ToolResult(
|
||
success=False,
|
||
error=str(e)
|
||
)
|
||
finally:
|
||
conn.close()
|
||
|
||
async def list_jobs(self) -> ToolResult:
|
||
"""Получить список всех задач."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute('''
|
||
SELECT id, name, command, schedule, enabled, last_run, next_run, created_at
|
||
FROM cron_jobs ORDER BY id
|
||
''')
|
||
rows = c.fetchall()
|
||
conn.close()
|
||
|
||
jobs = []
|
||
for row in rows:
|
||
jobs.append({
|
||
'id': row[0],
|
||
'name': row[1],
|
||
'command': row[2],
|
||
'schedule': row[3],
|
||
'enabled': bool(row[4]),
|
||
'last_run': row[5],
|
||
'next_run': row[6],
|
||
'created_at': row[7]
|
||
})
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data=jobs,
|
||
metadata={'count': len(jobs)}
|
||
)
|
||
|
||
async def remove_job(self, job_id: int) -> ToolResult:
|
||
"""Удалить задачу."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute("DELETE FROM cron_jobs WHERE id = ?", (job_id,))
|
||
if c.rowcount == 0:
|
||
conn.close()
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"Задача не найдена: {job_id}"
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
if job_id in self._jobs:
|
||
del self._jobs[job_id]
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data={'id': job_id},
|
||
metadata={'status': 'removed'}
|
||
)
|
||
|
||
async def toggle_job(self, job_id: int, enabled: bool) -> ToolResult:
|
||
"""Включить/выключить задачу."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute("UPDATE cron_jobs SET enabled = ? WHERE id = ?", (1 if enabled else 0, job_id))
|
||
if c.rowcount == 0:
|
||
conn.close()
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"Задача не найдена: {job_id}"
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data={'id': job_id, 'enabled': enabled},
|
||
metadata={'status': 'toggled'}
|
||
)
|
||
|
||
async def run_job(self, job_id: int) -> ToolResult:
|
||
"""Выполнить задачу немедленно."""
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute("SELECT command FROM cron_jobs WHERE id = ?", (job_id,))
|
||
row = c.fetchone()
|
||
|
||
if not row:
|
||
conn.close()
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"Задача не найдена: {job_id}"
|
||
)
|
||
|
||
command = row[0]
|
||
conn.close()
|
||
|
||
# Здесь должна быть логика выполнения команды
|
||
# Для демонстрации возвращаем заглушку
|
||
logger.info(f"Выполнение задачи {job_id}: {command}")
|
||
|
||
# Обновляем last_run
|
||
conn = sqlite3.connect(self.db_path)
|
||
c = conn.cursor()
|
||
c.execute("UPDATE cron_jobs SET last_run = datetime('now') WHERE id = ?", (job_id,))
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
return ToolResult(
|
||
success=True,
|
||
data={'id': job_id, 'command': command, 'message': 'Задача выполнена'},
|
||
metadata={'status': 'executed'}
|
||
)
|
||
|
||
async def execute(self, action: str = "list", **kwargs) -> ToolResult:
|
||
"""
|
||
Выполнить действие с cron задачами.
|
||
|
||
Args:
|
||
action: Действие - list, add, remove, toggle, run
|
||
kwargs: Дополнительные аргументы
|
||
"""
|
||
actions = {
|
||
'list': self.list_jobs,
|
||
'add': lambda: self.add_job(
|
||
name=kwargs.get('name'),
|
||
command=kwargs.get('command'),
|
||
schedule=kwargs.get('schedule')
|
||
),
|
||
'remove': lambda: self.remove_job(job_id=kwargs.get('job_id')),
|
||
'toggle': lambda: self.toggle_job(
|
||
job_id=kwargs.get('job_id'),
|
||
enabled=kwargs.get('enabled', True)
|
||
),
|
||
'run': lambda: self.run_job(job_id=kwargs.get('job_id'))
|
||
}
|
||
|
||
if action not in actions:
|
||
return ToolResult(
|
||
success=False,
|
||
error=f"Неизвестное действие: {action}. Доступные: {list(actions.keys())}"
|
||
)
|
||
|
||
logger.info(f"Cron действие: {action} с аргументами: {kwargs}")
|
||
return await actions[action]()
|
||
|
||
|
||
# Автоматическая регистрация при импорте
|
||
@register_tool
|
||
class CronToolAuto(CronTool):
|
||
"""Авто-регистрируемая версия CronTool."""
|
||
pass
|