#!/usr/bin/env/python
#_*_coding:utf-8_*_
#Data:17-10-08
#Auther:蘇莫
#Link:http://blog.csdn.net/lingluofengzang
#PythonVersion:python2.7
#filename:download_movie.py

import os
import sys
import requests

reload(sys)
sys.setdefaultencoding('utf-8')

# 功能：爬取m3u8格式的視頻

# 檢查存儲路徑是否正常
def check_path(_path):
  # 判斷存儲路徑是否存在
  if os.path.isdir(_path) or os.path.isabs(_path):
    # 判斷存儲路徑是否為空
    if not os.listdir(_path):
      return _path

    else:

      print u'>>>[-] 目標文件不為空，將清空目標文件，是否更換路徑？'
      flag = raw_input('>>>[*] Yes:1 No:2 \n>>>[+] [2]')

      try:
        if flag == '1':
          _path = raw_input(unicode('>>>[+] 請輸入目標文件路徑。\n>>>[+] ').encode('gbk'))
          check_path(_path)
        else:
          # 清空存儲路徑
          os.system('rd /S /Q ' + _path)
          os.system('mkdir ' + _path)
          return _path
      except Exception as e:
        print e
        exit(0)

  else:
    os.makedirs(_path)
    return _path

# 獲取ts視頻的爬取位置
def get_url(_url, _path):

  all_url = _url.split('/')
  url_pre = '/'.join(all_url[:-1]) + '/'
  url_next = all_url[-1]

  os.chdir(_path)
  # 獲取m3u8文件
  m3u8_txt = requests.get(_url, headers = {'Connection':'close'})
  with open(url_next, 'wb') as m3u8_content:
    m3u8_content.write(m3u8_txt.content)
  # 提取ts視頻的url
  movies_url = []
  _urls = open(url_next, 'rb')
  for line in _urls.readlines():
    if '.ts' in line:
      movies_url.append(url_pre + line[:-1])
    else:
      continue

  _urls.close()
  return movies_url

# 爬取ts視頻
def download_movie(movie_url, _path):
  os.chdir(_path)
  print '>>>[+] downloading...'
  print '-' * 60
  error_get = []

  for _url in movie_url:
    # ts視頻的名稱
    movie_name = _url.split('/')[-1][-6:]

    try:
      # 'Connection':'close' 防止請求端口占用
      # timeout=30  防止請求時間超長連接
      movie = requests.get(_url, headers = {'Connection':'close'}, timeout=60)
      with open(movie_name, 'wb') as movie_content:
        movie_content.writelines(movie)
      print '>>>[+] File ' + movie_name + ' done'
    # 捕獲異常，記錄失敗請求
    except:
      error_get.append(_url)
      continue
  # 如果沒有不成功的請求就結束
  if error_get:
    print u'共有%d個請求失敗' % len(file_list)
    print '-' * 60
    download_movie(error_get, _path)
  else:
    print '>>>[+] Download successfully!!!'

if __name__ == '__main__':
  try:

    _url = raw_input(unicode('>>>[+] 請輸入指定的[.m3u8]目標URL。\n>>>[+] ').encode('gbk'))
    _path = raw_input(unicode('>>>[+] 請輸入存儲目標文件路徑。\n>>>[+] ').encode('gbk'))

    storage_path = check_path(_path)
    movie_url = get_url(_url, storage_path)
    download_movie(movie_url, storage_path)

  except Exception as e:
    print e